diff --git a/personalised/code/configs/README.md b/personalised/code/configs/README.md deleted file mode 100644 index cf16b6458d7461b93c15e44353b1cde2fad05975..0000000000000000000000000000000000000000 --- a/personalised/code/configs/README.md +++ /dev/null @@ -1,27 +0,0 @@ -# Config Organization - -The task-facing Hydra configs are organized by challenge setting. Use these files with `python main.py --config-name
/ ...`. - -| Section | Entry config | Training / evaluation entry | Model config files | Model code | Checkpoints | -| --- | --- | --- | --- | --- | --- | -| Generic online | `generic_online/motion_diffusion.yaml` | `main.py` with `stage=fit` or `stage=test` | `generic_online/model/motion_diffusion.yaml`, `generic_online/model/losses/motion_diffusion.yaml` | `trainer/motion_diffusion.py`, `framework/motion_diffusion/` | `save/motion_diffusion/react_2025/online/checkpoints//` | -| Generic online | `generic_online/motion_transvae.yaml` | `main.py` with `stage=fit` or `stage=test` | `generic_online/model/motion_transvae.yaml`, `generic_online/model/losses/motion_transvae.yaml` | `trainer/motion_transvae.py`, `framework/motion_transvae/` | `save/motion_transvae/react_2025/online/checkpoints//` | -| Generic offline | `generic_offline/motion_diffusion.yaml` | `main.py` with `stage=fit` or `stage=test` | `generic_offline/model/motion_diffusion.yaml`, `generic_offline/model/losses/motion_diffusion.yaml` | `trainer/motion_diffusion.py`, `framework/motion_diffusion/` | `save/motion_diffusion/react_2025/offline/checkpoints//` | -| Generic offline | `generic_offline/motion_transvae.yaml` | `main.py` with `stage=fit` or `stage=test` | `generic_offline/model/motion_transvae.yaml`, `generic_offline/model/losses/motion_transvae.yaml` | `trainer/motion_transvae.py`, `framework/motion_transvae/` | `save/motion_transvae/react_2025/offline/checkpoints//` | -| Generic offline | RegNN | `regnn/train.py`; add `--test` for evaluation | RegNN uses command-line arguments instead of Hydra YAML | `regnn/models/`, `regnn/trainers.py` | `regnn//mhp-*-seed.pth` or `regnn//mhp-eeg-head-*-seed.pth` | -| Personalized online | `personalized_online/perfrdiff_rewrite_weight.yaml` | `main.py` with `stage=fit` or `stage=test` | `personalized_online/model/motion_diffusion.yaml`, `personalized_online/model/losses/perfrdiff_rewrite_weight.yaml` | `trainer/perfrdiff_rewrite_weight.py`, `framework/perfrdiff_rewrite_weight/` | `save/perfrdiff_rewrite_weight/react_2025/online/checkpoints//` | -| Personalized offline | `personalized_offline/perfrdiff_rewrite_weight.yaml` | `main.py` with `stage=fit` or `stage=test` | `personalized_offline/model/motion_diffusion.yaml`, `personalized_offline/model/losses/perfrdiff_rewrite_weight.yaml` | `trainer/perfrdiff_rewrite_weight.py`, `framework/perfrdiff_rewrite_weight/` | `save/perfrdiff_rewrite_weight/react_2025/offline/checkpoints//` | - -## Directory Notes - -- `data/`, `trainer/`, `model/`, and `model/losses/` under each section contain the task-specific YAML used by that section. -- `shared/` contains global support configs: `data/`, `logger/`, `model/`, `trainer/`, and `path.yaml`. -- `configs/shared/path.yaml` exposes resolver-based paths such as the code root and current working directory, so it is kept with the other shared configs. -- `configs/shared/model/emotion_autoencoder.yaml` is intentionally kept as a shared config because the post-processor loads it by this path. - -## Running - -- Hydra training: `python main.py --config-name
/ stage=fit data_dir=` -- Hydra evaluation: `python main.py --config-name
/ stage=test data_dir= resume_id=` -- RegNN training: run `python train.py ...` from the `regnn/` directory. -- RegNN evaluation: run `python train.py --test ...` from the `regnn/` directory. diff --git a/personalised/code/configs/data/emotion_autoencoder.yaml b/personalised/code/configs/data/emotion_autoencoder.yaml deleted file mode 100644 index c1578cc203b43cf4f693650cea47a429b13937df..0000000000000000000000000000000000000000 --- a/personalised/code/configs/data/emotion_autoencoder.yaml +++ /dev/null @@ -1,29 +0,0 @@ -data_name: react_2025 -_target_: dataset.react_2025.ReactionDataloader - -train_dataset: - _target_: dataset.react_2025.ReactionAutoEncoderDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 2 - root_dir: ${data_dir} - split: train - clip_length: ${trainer.clip_length} - -validation_dataset: - _target_: dataset.react_2025.ReactionAutoEncoderDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 2 - root_dir: ${data_dir} - split: val - clip_length: ${trainer.clip_length} - -test_dataset: - _target_: dataset.react_2025.ReactionAutoEncoderDataset - batch_size: 1 - shuffle: False - num_workers: 2 - root_dir: ${data_dir} - split: test - clip_length: ${trainer.clip_length} \ No newline at end of file diff --git a/personalised/code/configs/data/motion_diffusion.yaml b/personalised/code/configs/data/motion_diffusion.yaml deleted file mode 100644 index 05d60f88d8974e59c20120be88a05c0d84d4fc4e..0000000000000000000000000000000000000000 --- a/personalised/code/configs/data/motion_diffusion.yaml +++ /dev/null @@ -1,122 +0,0 @@ -data_name: react_2025 -_target_: dataset.react_2025.ReactionDataloader - -train_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 8 - root_dir: ${data_dir} - split: train - clip_length: 750 # 256 (transvae) - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - -validation_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 8 - root_dir: ${data_dir} - split: val - target_size: 224 - crop_size: 224 - clip_length: 750 # 256 (transvae) - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - -test_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 4 - root_dir: ${data_dir} - split: test - target_size: 224 - crop_size: 224 - clip_length: 750 # 256 (transvae) - fps: 30 - audio_feature_type: wav2vec # mfcc | wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: ${trainer.generic.eval_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} diff --git a/personalised/code/configs/data/motion_transvae.yaml b/personalised/code/configs/data/motion_transvae.yaml deleted file mode 100644 index b6b2c154a346926125a29a335450fb6b2a5193c6..0000000000000000000000000000000000000000 --- a/personalised/code/configs/data/motion_transvae.yaml +++ /dev/null @@ -1,125 +0,0 @@ -data_name: react_2025 -_target_: dataset.react_2025.ReactionDataloader - -train_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 8 - root_dir: ${data_dir} - split: train - clip_length: ${trainer.max_seq_len} - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: True - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.bidirectional} - normalize_3dmm: ${trainer.data_transform} - -validation_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 8 - root_dir: ${data_dir} - split: val - target_size: 224 - crop_size: 224 - clip_length: ${trainer.max_seq_len} - fps: 30 - audio_feature_type: wav2vec - load_video_s: True - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.bidirectional} - normalize_3dmm: ${trainer.data_transform} - -test_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: 1 - shuffle: True - num_workers: 0 - root_dir: ${data_dir} - split: test - target_size: 224 - crop_size: 224 - clip_length: ${trainer.max_seq_len} - fps: 30 - audio_feature_type: wav2vec # mfcc | wav2vec - load_video_s: True - load_video_l: ${trainer.renderer.do_render} - load_audio: True - load_emotion_s: True - load_emotion_l: ${trainer.eval_facial_metrics} - load_3dmm_s: True - load_3dmm_l: False - load_eeg_l: ${trainer.eval_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.bidirectional} - normalize_3dmm: ${trainer.data_transform} diff --git a/personalised/code/configs/data/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/data/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index 06e8e5fe9131535e98886a7e3f694dc46023a1c8..0000000000000000000000000000000000000000 --- a/personalised/code/configs/data/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,135 +0,0 @@ -data_name: react_2025 -_target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataModule - -train_dataset: - _target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 8 - root_dir: ${data_dir} - split: train - clip_length: 750 - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_ref: True - load_personality_l: True - personal_condition_mode: ${trainer.main_model.args.personal_condition_mode} - load_eeg_l: ${trainer.generic.train_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - normalize_3dmm: standard - -validation_dataset: - _target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 8 - root_dir: ${data_dir} - split: val - clip_length: 750 - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_ref: True - load_personality_l: True - personal_condition_mode: ${trainer.main_model.args.personal_condition_mode} - load_eeg_l: ${trainer.generic.train_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - normalize_3dmm: standard - -test_dataset: - _target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 4 - root_dir: ${data_dir} - split: test - clip_length: 750 - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_ref: False - load_personality_l: True - personal_condition_mode: ${trainer.main_model.args.personal_condition_mode} - load_eeg_l: ${trainer.generic.eval_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - normalize_3dmm: standard - num_test_gts: ${num_gts} diff --git a/personalised/code/configs/g2p_delta.yaml b/personalised/code/configs/g2p_delta.yaml deleted file mode 100644 index 220fef845c2016a301e462bc5639386abf7fa0ea..0000000000000000000000000000000000000000 --- a/personalised/code/configs/g2p_delta.yaml +++ /dev/null @@ -1,88 +0,0 @@ -defaults: - - /personalized_offline/data/perfrdiff_rewrite_weight@data - - /personalized_offline/trainer/perfrdiff_rewrite_weight@trainer - - /shared/logger: none - - /shared/path@path - - _self_ - -seed: 1234 -logger_level: INFO -task: offline -stage: ??? -data_dir: ??? -resume: false -run_id: ${generate_id:} -resume_id: "" -num_gts: 10 - -trainer: - _target_: trainer.g2p_delta.Trainer - model_name: g2p_delta - batch_size: 2 - pretrained: - diffusion_decoder: pretrained_models/generic_offline_final_submission_model/CausalTransformerDenoiser/checkpoint_120.pth - diffusion_prior: pretrained_models/generic_offline_final_submission_model/DiffusionPriorNetwork/checkpoint_120.pth - modifier_checkpoint: "" - eeg_head_checkpoint: pretrained_models/generic_offline_final_submission_model/EEGPredictionHead/checkpoint_120.pth - model: - diff_model: - _target_: framework.motion_diffusion.diffusion.matchers_causal.CausalLatentMatcher - eeg_head: - enabled: true - diffusion_decoder: - args: - use_lag_bias: true - lag_max: 60 - lag_lookahead: 0 - use_coarse: true - coarse_classes: 8 - coarse_hidden: 256 - coarse_emo_start: 17 - criterion: - args: - eeg_loss_weight: 0.25 - generic: - epochs: 30 - save_period: 5 - val_period: 5 - clip_grad: true - train_eeg: true - train_eeg_head_only: false - eval_eeg: true - num_workers: 8 - num_preds: 1 - max_train_batches: 0 - max_val_batches: 0 - # Short-triage control only: matched (default) | shuffled | identity. - # See trainer/g2p_delta.py::_apply_personalization and - # scripts/build_subset_eval_root.py. Never used during training. - eval_condition_mode: matched - main_model: - args: - personal_condition_mode: personality_only - embed_dim: 256 - history_hidden_dim: 128 - personality_hidden_dim: 128 - condition_dropout: 0.1 - delta_rank: 128 - delta_max_scale: 0.15 - personalize_coarse: true - anchor_weight: 0.001 - coarse_weight: 0.5 - ccc_weight: 0.1 - dynamics_weight: 0.02 - counterfactual_weight: 0.1 - counterfactual_margin: 0.02 - eeg_lr: 0.00001 - optimizer_hypernet: - type: adamw - args: - lr: 0.0002 - weight_decay: 0.0001 - momentum: 0.0 - -hydra: - run: - dir: outputs/${trainer.model_name}/${data.data_name}/${task}/${run_id} - job: - chdir: true diff --git a/personalised/code/configs/generic_offline/data/motion_diffusion.yaml b/personalised/code/configs/generic_offline/data/motion_diffusion.yaml deleted file mode 100644 index 05d60f88d8974e59c20120be88a05c0d84d4fc4e..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/data/motion_diffusion.yaml +++ /dev/null @@ -1,122 +0,0 @@ -data_name: react_2025 -_target_: dataset.react_2025.ReactionDataloader - -train_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 8 - root_dir: ${data_dir} - split: train - clip_length: 750 # 256 (transvae) - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - -validation_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 8 - root_dir: ${data_dir} - split: val - target_size: 224 - crop_size: 224 - clip_length: 750 # 256 (transvae) - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - -test_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 4 - root_dir: ${data_dir} - split: test - target_size: 224 - crop_size: 224 - clip_length: 750 # 256 (transvae) - fps: 30 - audio_feature_type: wav2vec # mfcc | wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: ${trainer.generic.eval_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} diff --git a/personalised/code/configs/generic_offline/data/motion_transvae.yaml b/personalised/code/configs/generic_offline/data/motion_transvae.yaml deleted file mode 100644 index b6b2c154a346926125a29a335450fb6b2a5193c6..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/data/motion_transvae.yaml +++ /dev/null @@ -1,125 +0,0 @@ -data_name: react_2025 -_target_: dataset.react_2025.ReactionDataloader - -train_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 8 - root_dir: ${data_dir} - split: train - clip_length: ${trainer.max_seq_len} - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: True - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.bidirectional} - normalize_3dmm: ${trainer.data_transform} - -validation_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 8 - root_dir: ${data_dir} - split: val - target_size: 224 - crop_size: 224 - clip_length: ${trainer.max_seq_len} - fps: 30 - audio_feature_type: wav2vec - load_video_s: True - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.bidirectional} - normalize_3dmm: ${trainer.data_transform} - -test_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: 1 - shuffle: True - num_workers: 0 - root_dir: ${data_dir} - split: test - target_size: 224 - crop_size: 224 - clip_length: ${trainer.max_seq_len} - fps: 30 - audio_feature_type: wav2vec # mfcc | wav2vec - load_video_s: True - load_video_l: ${trainer.renderer.do_render} - load_audio: True - load_emotion_s: True - load_emotion_l: ${trainer.eval_facial_metrics} - load_3dmm_s: True - load_3dmm_l: False - load_eeg_l: ${trainer.eval_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.bidirectional} - normalize_3dmm: ${trainer.data_transform} diff --git a/personalised/code/configs/generic_offline/model/losses/motion_diffusion.yaml b/personalised/code/configs/generic_offline/model/losses/motion_diffusion.yaml deleted file mode 100644 index f528b1e1cb4be34a38fde02543cc3945f2dd520f..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/model/losses/motion_diffusion.yaml +++ /dev/null @@ -1,11 +0,0 @@ -_target_: framework.utils.losses.DiffusionLoss - -losses_type: MSELossApt -losses_multiplier: 1.0 -n_preds: ${trainer.${trainer.trainer_mode}.num_preds} -temporal_loss_w: 0.0 -prior_loss_weight: 1.0 -eeg_loss_weight: 0.25 -w_au: 1.0 # action unit -w_va: 5.0 # valence and arousal -w_em: 5.0 # emotion diff --git a/personalised/code/configs/generic_offline/model/losses/motion_transvae.yaml b/personalised/code/configs/generic_offline/model/losses/motion_transvae.yaml deleted file mode 100644 index fd96f6674d5a4b7ddfc2caf12a95c19b4ef4d7ea..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/model/losses/motion_transvae.yaml +++ /dev/null @@ -1,8 +0,0 @@ -_target_: framework.utils.losses.VAELoss - -kl_p: 0.00001 # kl div loss -w_emo: 2.0 -w_exp: 2.0 # facial expression -w_rot: 4.0 # rotation -w_tran: 4.0 # translation -eeg_loss_weight: 0.25 diff --git a/personalised/code/configs/generic_offline/model/motion_diffusion.yaml b/personalised/code/configs/generic_offline/model/motion_diffusion.yaml deleted file mode 100644 index 59385f65e8a2d603bc2c32b5abea0e8ee481cc02..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/model/motion_diffusion.yaml +++ /dev/null @@ -1,116 +0,0 @@ -diff_model: - model_name: LatentMLPMatcher - _target_: framework.motion_diffusion.diffusion.matchers.LatentMatcher - - eeg_head: - enabled: True - output_dim: 14 - hidden_dim: 256 - dropout: 0.5 - pooling: mean - detach_prediction_emotion: True - use_speaker_audio: True - use_speaker_emotion: True - use_speaker_3dmm: True - use_prediction_emotion: True - - diffusion_prior: - type: - DiffusionPriorNetwork - args: - emb_preprocessing: normalize - freeze_encoder: True - audio_dim: 768 - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - _3dmm_dim: 58 - speaker_emb_dim: 512 - latent_dim: 512 - depth: 4 - num_time_layers: 2 - num_time_embeds: 1 - num_time_emb_channels: 64 - time_last_act: False - use_learned_query: True - s_audio_cond_drop_prob: 0.2 - s_latentemb_cond_drop_prob: 1.0 - s_3dmm_cond_drop_prob: 0.2 - guidance_scale: 1.0 - dim_head: 64 - heads: 8 - ff_mult: 4 - norm_in: False - norm_out: True - attn_dropout: 0.0 - ff_dropout: 0.0 - final_proj: True - normformer: False - rotary_emb: True - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - - diffusion_decoder: - type: - TransformerDenoiser - args: - emb_preprocessing: normalize - freeze_encoder: True - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - encode_emotion: False - encode_3dmm: False - ablation_skip_connection: True - nfeats: 25 - latent_dim: 512 - ff_size: 1024 - num_layers: 9 # 7 | 9 - num_heads: 8 # 4 | 8 - dropout: 0.1 - normalize_before: False - activation: gelu - flip_sin_to_cos: True - return_intermediate_dec: False - position_embedding: learned - arch: trans_dec - freq_shift: 0 - time_encoded_dim: 64 - s_audio_dim: 768 - s_audio_scale: 1.0 - s_emotion_dim: 25 - l_embed_dim: 512 - s_embed_dim: 512 - personal_emb_dim: 512 - s_3dmm_dim: 58 - concat: concat_first - guidance_scale: 1.0 # 7.5 - s_audio_enc_drop_prob: 0.2 - s_latent_embed_drop_prob: 1.0 - s_3dmm_enc_drop_prob: 0.2 - s_emotion_enc_drop_prob: 0.2 - past_l_emotion_drop_prob: 1.0 - use_past_frames: False - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - -defaults: - - /generic_offline/model/losses/motion_diffusion@loss - - /generic_offline/model/motion_diffusion/audio_embedder@audio_encoder - #- /model/latent_embedder@latent_embedder diff --git a/personalised/code/configs/generic_offline/model/motion_diffusion/audio_embedder.yaml b/personalised/code/configs/generic_offline/model/motion_diffusion/audio_embedder.yaml deleted file mode 100644 index 8e18655a404101b4ee85faa3956a14c00ea1cc4b..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/model/motion_diffusion/audio_embedder.yaml +++ /dev/null @@ -1,4 +0,0 @@ -type: AudioEmbedder -_target_: dataset.modules.audio_embedder.AudioEmbedder -checkpoint_path: "" -skip_norm: True \ No newline at end of file diff --git a/personalised/code/configs/generic_offline/model/motion_transvae.yaml b/personalised/code/configs/generic_offline/model/motion_transvae.yaml deleted file mode 100644 index 08e0afc25b4cadc9af577a8faefb0761cdd8f722..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/model/motion_transvae.yaml +++ /dev/null @@ -1,24 +0,0 @@ -model_name: TransformerVAE -_target_: framework.motion_transvae.TransformerVAE.TransformerVAE - -img_size: 224 -audio_dim: 768 -output_3dmm_dim: 58 -output_emotion_dim: 25 -feature_dim: 128 -seq_len: ${data.train_dataset.clip_length} -task: ${task} -window_size: ${trainer.window_size} -device: cuda - -eeg_head: - enabled: True - output_dim: 14 - hidden_dim: 256 - dropout: 0.5 - pooling: mean - detach_prediction_emotion: True - use_speaker_audio: True - use_speaker_emotion: True - use_speaker_3dmm: True - use_prediction_emotion: True diff --git a/personalised/code/configs/generic_offline/motion_diffusion.yaml b/personalised/code/configs/generic_offline/motion_diffusion.yaml deleted file mode 100644 index 61651ade05e88d28af3ffa4e65d7618fc49fde41..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/motion_diffusion.yaml +++ /dev/null @@ -1,23 +0,0 @@ -hydra: - run: - dir: outputs/${trainer.model_name}/${data.data_name}/${task}/${run_id} - job: - chdir: true - -seed: 1234 -logger_level: INFO - -task: offline -stage: ??? -data_dir: ??? -resume: false -run_id: ${generate_id:} -resume_id: "" -num_gts: 10 - -defaults: - - /generic_offline/data/motion_diffusion@data - - /generic_offline/trainer/motion_diffusion@trainer - - /shared/logger: none - - /shared/path@path - - _self_ diff --git a/personalised/code/configs/generic_offline/motion_transvae.yaml b/personalised/code/configs/generic_offline/motion_transvae.yaml deleted file mode 100644 index 47f14222e2a4fd1fad4ec94afb22b927cb900cf8..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/motion_transvae.yaml +++ /dev/null @@ -1,23 +0,0 @@ -hydra: - run: - dir: outputs/${trainer.model_name}/${data.data_name}/${task}/${run_id} - job: - chdir: true - -seed: 1234 -logger_level: INFO - -task: offline -stage: ??? -data_dir: ??? -resume: false -run_id: ${generate_id:} -resume_id: "" -num_gts: 10 - -defaults: - - /generic_offline/data/motion_transvae@data - - /generic_offline/trainer/motion_transvae@trainer - - /shared/logger: none - - /shared/path@path - - _self_ diff --git a/personalised/code/configs/generic_offline/trainer/motion_diffusion.yaml b/personalised/code/configs/generic_offline/trainer/motion_diffusion.yaml deleted file mode 100644 index b2ca1d565e81819eaeac2485c2b2bf3f7c10bf87..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/trainer/motion_diffusion.yaml +++ /dev/null @@ -1,48 +0,0 @@ -model_name: motion_diffusion -_target_: trainer.motion_diffusion.Trainer - -trainer_mode: generic # generic facial reaction generator -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} # true | false -current_runid: ${run_id} -resume_runid: ${resume_id} -batch_size: 4 -post_config_name: configs/shared/model/emotion_autoencoder.yaml -post_clip_length: 1000 # ground-truths alignment -data_clamp: True - -generic: - seed: 1234 - start_epoch: 0 - epochs: 500 - lr: 0.0001 - model: LatentMatcher - clip_grad: False - num_workers: 16 - log_dir: log - tb_dir: tb_logs - out_dir: results - save_period: 10 - val_period: 10 - num_preds: 10 - clip_length: 750 - window_size: 30 - s_ratio: 2 - train_eeg_head_only: False - pretrained_decoder_checkpoint: "" - pretrained_prior_checkpoint: "" - pretrained_load_prior: True - eval_eeg: True - bidirectional: False - -renderer: # face animation based on predicted 3DMM coefficients - name: PIRenderer - _target_: utils.render.Render - do_render: False - -defaults: - - /generic_offline/model/motion_diffusion@model - - /shared/model/optim/adamw@optim - - /shared/model/scheduler/cosine_annealing@scheduler - - /generic_offline/model/losses/motion_diffusion@criterion # loss func diff --git a/personalised/code/configs/generic_offline/trainer/motion_transvae.yaml b/personalised/code/configs/generic_offline/trainer/motion_transvae.yaml deleted file mode 100644 index 2a7768d1c16df8375a0075a0619734cf0fe4ca03..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_offline/trainer/motion_transvae.yaml +++ /dev/null @@ -1,42 +0,0 @@ -model_name: motion_transvae -_target_: trainer.motion_transvae.Trainer - -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} # true | false -current_runid: ${run_id} -resume_runid: ${resume_id} - -epochs: 100 -gpu_ids: 0 -lr: 0.00001 -j: 12 # num_workers -batch_size: 4 -max_seq_len: 750 -window_size: 8 -div_p: 10 -post_config_name: configs/shared/model/emotion_autoencoder.yaml -post_clip_length: 1000 # ground-truths alignment -data_clamp: True -data_transform: zero_center # standard | zero_center -train_eeg_head_only: False -pretrained_model_checkpoint: "" -eval_eeg: True -num_preds: 10 -bidirectional: False -save_results: True -eval_facial_metrics: True -eval_eeg_metrics: True -metric_threads: 1 -eval_clip_batch_size: 1 - -renderer: # face animation based on predicted 3DMM coefficients - name: PIRenderer - _target_: utils.render.Render - do_render: True - transform_reverse: ${trainer.data_transform} - -defaults: - - /generic_offline/model/motion_transvae@model - - /shared/model/optim/adamw@optim - - /generic_offline/model/losses/motion_transvae@criterion # loss func diff --git a/personalised/code/configs/generic_online/data/motion_diffusion.yaml b/personalised/code/configs/generic_online/data/motion_diffusion.yaml deleted file mode 100644 index 05d60f88d8974e59c20120be88a05c0d84d4fc4e..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/data/motion_diffusion.yaml +++ /dev/null @@ -1,122 +0,0 @@ -data_name: react_2025 -_target_: dataset.react_2025.ReactionDataloader - -train_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 8 - root_dir: ${data_dir} - split: train - clip_length: 750 # 256 (transvae) - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - -validation_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 8 - root_dir: ${data_dir} - split: val - target_size: 224 - crop_size: 224 - clip_length: 750 # 256 (transvae) - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - -test_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 4 - root_dir: ${data_dir} - split: test - target_size: 224 - crop_size: 224 - clip_length: 750 # 256 (transvae) - fps: 30 - audio_feature_type: wav2vec # mfcc | wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: ${trainer.generic.eval_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} diff --git a/personalised/code/configs/generic_online/data/motion_transvae.yaml b/personalised/code/configs/generic_online/data/motion_transvae.yaml deleted file mode 100644 index b6b2c154a346926125a29a335450fb6b2a5193c6..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/data/motion_transvae.yaml +++ /dev/null @@ -1,125 +0,0 @@ -data_name: react_2025 -_target_: dataset.react_2025.ReactionDataloader - -train_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 8 - root_dir: ${data_dir} - split: train - clip_length: ${trainer.max_seq_len} - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: True - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.bidirectional} - normalize_3dmm: ${trainer.data_transform} - -validation_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 8 - root_dir: ${data_dir} - split: val - target_size: 224 - crop_size: 224 - clip_length: ${trainer.max_seq_len} - fps: 30 - audio_feature_type: wav2vec - load_video_s: True - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_eeg_l: True - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.bidirectional} - normalize_3dmm: ${trainer.data_transform} - -test_dataset: - _target_: dataset.react_2025.ReactionDataset - batch_size: 1 - shuffle: True - num_workers: 0 - root_dir: ${data_dir} - split: test - target_size: 224 - crop_size: 224 - clip_length: ${trainer.max_seq_len} - fps: 30 - audio_feature_type: wav2vec # mfcc | wav2vec - load_video_s: True - load_video_l: ${trainer.renderer.do_render} - load_audio: True - load_emotion_s: True - load_emotion_l: ${trainer.eval_facial_metrics} - load_3dmm_s: True - load_3dmm_l: False - load_eeg_l: ${trainer.eval_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.bidirectional} - normalize_3dmm: ${trainer.data_transform} diff --git a/personalised/code/configs/generic_online/model/losses/motion_diffusion.yaml b/personalised/code/configs/generic_online/model/losses/motion_diffusion.yaml deleted file mode 100644 index f528b1e1cb4be34a38fde02543cc3945f2dd520f..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/model/losses/motion_diffusion.yaml +++ /dev/null @@ -1,11 +0,0 @@ -_target_: framework.utils.losses.DiffusionLoss - -losses_type: MSELossApt -losses_multiplier: 1.0 -n_preds: ${trainer.${trainer.trainer_mode}.num_preds} -temporal_loss_w: 0.0 -prior_loss_weight: 1.0 -eeg_loss_weight: 0.25 -w_au: 1.0 # action unit -w_va: 5.0 # valence and arousal -w_em: 5.0 # emotion diff --git a/personalised/code/configs/generic_online/model/losses/motion_transvae.yaml b/personalised/code/configs/generic_online/model/losses/motion_transvae.yaml deleted file mode 100644 index fd96f6674d5a4b7ddfc2caf12a95c19b4ef4d7ea..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/model/losses/motion_transvae.yaml +++ /dev/null @@ -1,8 +0,0 @@ -_target_: framework.utils.losses.VAELoss - -kl_p: 0.00001 # kl div loss -w_emo: 2.0 -w_exp: 2.0 # facial expression -w_rot: 4.0 # rotation -w_tran: 4.0 # translation -eeg_loss_weight: 0.25 diff --git a/personalised/code/configs/generic_online/model/motion_diffusion.yaml b/personalised/code/configs/generic_online/model/motion_diffusion.yaml deleted file mode 100644 index eb29cb006f990b7bcaa3a2a238aa8e0633771967..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/model/motion_diffusion.yaml +++ /dev/null @@ -1,116 +0,0 @@ -diff_model: - model_name: LatentMLPMatcher - _target_: framework.motion_diffusion.diffusion.matchers.LatentMatcher - - eeg_head: - enabled: True - output_dim: 14 - hidden_dim: 256 - dropout: 0.5 - pooling: mean - detach_prediction_emotion: True - use_speaker_audio: True - use_speaker_emotion: True - use_speaker_3dmm: True - use_prediction_emotion: True - - diffusion_prior: - type: - DiffusionPriorNetwork - args: - emb_preprocessing: normalize - freeze_encoder: True - audio_dim: 768 - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - _3dmm_dim: 58 - speaker_emb_dim: 512 - latent_dim: 512 - depth: 4 - num_time_layers: 2 - num_time_embeds: 1 - num_time_emb_channels: 64 - time_last_act: False - use_learned_query: True - s_audio_cond_drop_prob: 0.2 - s_latentemb_cond_drop_prob: 1.0 - s_3dmm_cond_drop_prob: 0.2 - guidance_scale: 1.0 - dim_head: 64 - heads: 8 - ff_mult: 4 - norm_in: False - norm_out: True - attn_dropout: 0.0 - ff_dropout: 0.0 - final_proj: True - normformer: False - rotary_emb: True - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - - diffusion_decoder: - type: - TransformerDenoiser - args: - emb_preprocessing: normalize - freeze_encoder: True - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - encode_emotion: False - encode_3dmm: False - ablation_skip_connection: True - nfeats: 25 - latent_dim: 512 - ff_size: 1024 - num_layers: 9 # 7 | 9 - num_heads: 8 # 4 | 8 - dropout: 0.1 - normalize_before: False - activation: gelu - flip_sin_to_cos: True - return_intermediate_dec: False - position_embedding: learned - arch: trans_dec - freq_shift: 0 - time_encoded_dim: 64 - s_audio_dim: 768 - s_audio_scale: 1.0 - s_emotion_dim: 25 - l_embed_dim: 512 - s_embed_dim: 512 - personal_emb_dim: 512 - s_3dmm_dim: 58 - concat: concat_first - guidance_scale: 1.0 # 7.5 - s_audio_enc_drop_prob: 0.2 - s_latent_embed_drop_prob: 1.0 - s_3dmm_enc_drop_prob: 0.2 - s_emotion_enc_drop_prob: 0.2 - past_l_emotion_drop_prob: 1.0 - use_past_frames: False - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - -defaults: - - /generic_online/model/losses/motion_diffusion@loss - - /generic_online/model/motion_diffusion/audio_embedder@audio_encoder - #- /model/latent_embedder@latent_embedder diff --git a/personalised/code/configs/generic_online/model/motion_diffusion/audio_embedder.yaml b/personalised/code/configs/generic_online/model/motion_diffusion/audio_embedder.yaml deleted file mode 100644 index 8e18655a404101b4ee85faa3956a14c00ea1cc4b..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/model/motion_diffusion/audio_embedder.yaml +++ /dev/null @@ -1,4 +0,0 @@ -type: AudioEmbedder -_target_: dataset.modules.audio_embedder.AudioEmbedder -checkpoint_path: "" -skip_norm: True \ No newline at end of file diff --git a/personalised/code/configs/generic_online/model/motion_transvae.yaml b/personalised/code/configs/generic_online/model/motion_transvae.yaml deleted file mode 100644 index 08e0afc25b4cadc9af577a8faefb0761cdd8f722..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/model/motion_transvae.yaml +++ /dev/null @@ -1,24 +0,0 @@ -model_name: TransformerVAE -_target_: framework.motion_transvae.TransformerVAE.TransformerVAE - -img_size: 224 -audio_dim: 768 -output_3dmm_dim: 58 -output_emotion_dim: 25 -feature_dim: 128 -seq_len: ${data.train_dataset.clip_length} -task: ${task} -window_size: ${trainer.window_size} -device: cuda - -eeg_head: - enabled: True - output_dim: 14 - hidden_dim: 256 - dropout: 0.5 - pooling: mean - detach_prediction_emotion: True - use_speaker_audio: True - use_speaker_emotion: True - use_speaker_3dmm: True - use_prediction_emotion: True diff --git a/personalised/code/configs/generic_online/motion_diffusion.yaml b/personalised/code/configs/generic_online/motion_diffusion.yaml deleted file mode 100644 index 0d2d17d2d32413cb9c153824b5fee5ccdd074dd9..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/motion_diffusion.yaml +++ /dev/null @@ -1,23 +0,0 @@ -hydra: - run: - dir: outputs/${trainer.model_name}/${data.data_name}/${task}/${run_id} - job: - chdir: true - -seed: 1234 -logger_level: INFO - -task: online -stage: ??? -data_dir: ??? -resume: false -run_id: ${generate_id:} -resume_id: "" -num_gts: 10 - -defaults: - - /generic_online/data/motion_diffusion@data - - /generic_online/trainer/motion_diffusion@trainer - - /shared/logger: none - - /shared/path@path - - _self_ diff --git a/personalised/code/configs/generic_online/motion_transvae.yaml b/personalised/code/configs/generic_online/motion_transvae.yaml deleted file mode 100644 index dc6398335044355b2311881bff216d4d42983784..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/motion_transvae.yaml +++ /dev/null @@ -1,23 +0,0 @@ -hydra: - run: - dir: outputs/${trainer.model_name}/${data.data_name}/${task}/${run_id} - job: - chdir: true - -seed: 1234 -logger_level: INFO - -task: online -stage: ??? -data_dir: ??? -resume: false -run_id: ${generate_id:} -resume_id: "" -num_gts: 10 - -defaults: - - /generic_online/data/motion_transvae@data - - /generic_online/trainer/motion_transvae@trainer - - /shared/logger: none - - /shared/path@path - - _self_ diff --git a/personalised/code/configs/generic_online/trainer/motion_diffusion.yaml b/personalised/code/configs/generic_online/trainer/motion_diffusion.yaml deleted file mode 100644 index 0712883a80805af1dde081226de910e6b64abb9f..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/trainer/motion_diffusion.yaml +++ /dev/null @@ -1,48 +0,0 @@ -model_name: motion_diffusion -_target_: trainer.motion_diffusion.Trainer - -trainer_mode: generic # generic facial reaction generator -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} # true | false -current_runid: ${run_id} -resume_runid: ${resume_id} -batch_size: 4 -post_config_name: configs/shared/model/emotion_autoencoder.yaml -post_clip_length: 1000 # ground-truths alignment -data_clamp: True - -generic: - seed: 1234 - start_epoch: 0 - epochs: 500 - lr: 0.0001 - model: LatentMatcher - clip_grad: False - num_workers: 16 - log_dir: log - tb_dir: tb_logs - out_dir: results - save_period: 10 - val_period: 10 - num_preds: 10 - clip_length: 750 - window_size: 30 - s_ratio: 2 - train_eeg_head_only: False - pretrained_decoder_checkpoint: "" - pretrained_prior_checkpoint: "" - pretrained_load_prior: True - eval_eeg: True - bidirectional: False - -renderer: # face animation based on predicted 3DMM coefficients - name: PIRenderer - _target_: utils.render.Render - do_render: False - -defaults: - - /generic_online/model/motion_diffusion@model - - /shared/model/optim/adamw@optim - - /shared/model/scheduler/cosine_annealing@scheduler - - /generic_online/model/losses/motion_diffusion@criterion # loss func diff --git a/personalised/code/configs/generic_online/trainer/motion_transvae.yaml b/personalised/code/configs/generic_online/trainer/motion_transvae.yaml deleted file mode 100644 index bcff8934156353f9474d99413ae2c81e23fbe640..0000000000000000000000000000000000000000 --- a/personalised/code/configs/generic_online/trainer/motion_transvae.yaml +++ /dev/null @@ -1,42 +0,0 @@ -model_name: motion_transvae -_target_: trainer.motion_transvae.Trainer - -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} # true | false -current_runid: ${run_id} -resume_runid: ${resume_id} - -epochs: 100 -gpu_ids: 0 -lr: 0.00001 -j: 12 # num_workers -batch_size: 4 -max_seq_len: 750 -window_size: 8 -div_p: 10 -post_config_name: configs/shared/model/emotion_autoencoder.yaml -post_clip_length: 1000 # ground-truths alignment -data_clamp: True -data_transform: zero_center # standard | zero_center -train_eeg_head_only: False -pretrained_model_checkpoint: "" -eval_eeg: True -num_preds: 10 -bidirectional: False -save_results: True -eval_facial_metrics: True -eval_eeg_metrics: True -metric_threads: 1 -eval_clip_batch_size: 1 - -renderer: # face animation based on predicted 3DMM coefficients - name: PIRenderer - _target_: utils.render.Render - do_render: True - transform_reverse: ${trainer.data_transform} - -defaults: - - /generic_online/model/motion_transvae@model - - /shared/model/optim/adamw@optim - - /generic_online/model/losses/motion_transvae@criterion # loss func diff --git a/personalised/code/configs/logger/none.yaml b/personalised/code/configs/logger/none.yaml deleted file mode 100644 index 3da2aede1c03e6d9660ccabe8ce187288aee3f48..0000000000000000000000000000000000000000 --- a/personalised/code/configs/logger/none.yaml +++ /dev/null @@ -1,4 +0,0 @@ -logger_name: none -version: ${run_id} -logger_dir: tb_logs -project: null \ No newline at end of file diff --git a/personalised/code/configs/main.yaml b/personalised/code/configs/main.yaml deleted file mode 100644 index 36e25fe9263b0dfb05a98eb37070e36e6f13e9db..0000000000000000000000000000000000000000 --- a/personalised/code/configs/main.yaml +++ /dev/null @@ -1,26 +0,0 @@ - -hydra: - run: - dir: outputs/${trainer.model_name}/${data.data_name}/${task}/${run_id} - job: # automatically go to the job folder (needed for hydra > 1.2 with new behavior) - chdir: true - -# Global configurations shared between different modules -seed: 1234 -logger_level: INFO - -task: ??? # online | offline (* mandatory) -stage: ??? # fit | test (* mandatory) -data_dir: ??? -# e.g., /lustre/projects/Research_Project-T127204/xk219/projects/datasets/react2026 -resume: false -run_id: ${generate_id:} # function defined in prepare.py -resume_id: "" # specified if resume -num_gts: 10 # ground-truths - -defaults: - - data: ??? # react_2025 (* mandatory) - - trainer: ??? # motion_diffusion | motion_transvae (* mandatory) - - logger: none - - /path@path # path configuration (loaded from configs/path.yaml) - - _self_ \ No newline at end of file diff --git a/personalised/code/configs/model/emotion_autoencoder.yaml b/personalised/code/configs/model/emotion_autoencoder.yaml deleted file mode 100644 index 2f80bbb7189d799d036a01b075ba3cc332a83b03..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/emotion_autoencoder.yaml +++ /dev/null @@ -1,17 +0,0 @@ -type: EmotionVAE -_target_: framework.modules.emotion_autoencoder.EmotionVAE - -in_channels: 25 -out_channels: 25 -feature_dim: 512 -nhead: 8 -dropout: 0.1 -num_encoder_layers: 6 -num_decoder_layers: 6 -mlp_dist: False # expand mu & logvar -in_proj_type : linear # linear | mlp -out_proj_type : separate # separate | shared -pe_type : absolute # learnable | absolute -query_type : zero -max_seq_len: 5000 -global_token_len: 64 \ No newline at end of file diff --git a/personalised/code/configs/model/losses/emotion_autoencoder.yaml b/personalised/code/configs/model/losses/emotion_autoencoder.yaml deleted file mode 100644 index 6881880215136f33e15ae18141d3470b69d6afc2..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/losses/emotion_autoencoder.yaml +++ /dev/null @@ -1,6 +0,0 @@ -_target_: framework.utils.losses.EmotionVAELoss - -w_au: 1.0 -w_va: 10.0 -w_em: 1.0 -w_kld: 0.001 \ No newline at end of file diff --git a/personalised/code/configs/model/losses/motion_diffusion.yaml b/personalised/code/configs/model/losses/motion_diffusion.yaml deleted file mode 100644 index f528b1e1cb4be34a38fde02543cc3945f2dd520f..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/losses/motion_diffusion.yaml +++ /dev/null @@ -1,11 +0,0 @@ -_target_: framework.utils.losses.DiffusionLoss - -losses_type: MSELossApt -losses_multiplier: 1.0 -n_preds: ${trainer.${trainer.trainer_mode}.num_preds} -temporal_loss_w: 0.0 -prior_loss_weight: 1.0 -eeg_loss_weight: 0.25 -w_au: 1.0 # action unit -w_va: 5.0 # valence and arousal -w_em: 5.0 # emotion diff --git a/personalised/code/configs/model/losses/motion_diffusion_causal.yaml b/personalised/code/configs/model/losses/motion_diffusion_causal.yaml deleted file mode 100644 index 9bd1b7aca31f445331118ddb566a2222d1fd3f31..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/losses/motion_diffusion_causal.yaml +++ /dev/null @@ -1,9 +0,0 @@ -# Loss with the explicit coarse-to-fine cross-entropy term added. -defaults: - - motion_diffusion - - _self_ - -_target_: framework.utils.losses_causal.DiffusionLossCoarse -w_coarse: 0.5 # weight of the coarse 8-class CE term -coarse_emo_start: 17 -coarse_classes: 8 diff --git a/personalised/code/configs/model/losses/motion_transvae.yaml b/personalised/code/configs/model/losses/motion_transvae.yaml deleted file mode 100644 index fd96f6674d5a4b7ddfc2caf12a95c19b4ef4d7ea..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/losses/motion_transvae.yaml +++ /dev/null @@ -1,8 +0,0 @@ -_target_: framework.utils.losses.VAELoss - -kl_p: 0.00001 # kl div loss -w_emo: 2.0 -w_exp: 2.0 # facial expression -w_rot: 4.0 # rotation -w_tran: 4.0 # translation -eeg_loss_weight: 0.25 diff --git a/personalised/code/configs/model/losses/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/model/losses/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index a8cd2da396773e073c381619488d0031e429476a..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/losses/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,8 +0,0 @@ -type: DiffusionLoss -args: - losses_type: [MSELoss, MSELoss] - losses_multipliers: [0, 1] - losses_decoded: [False, True] - k: ${trainer.${trainer.trainer_mode}.num_preds} - temporal_loss_w: 0.0 - eeg_loss_weight: 1.0 diff --git a/personalised/code/configs/model/motion_diffusion.yaml b/personalised/code/configs/model/motion_diffusion.yaml deleted file mode 100644 index fd313fcd0a1c96e3bb90552b6a3b51b14c336e05..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/motion_diffusion.yaml +++ /dev/null @@ -1,116 +0,0 @@ -diff_model: - model_name: LatentMLPMatcher - _target_: framework.motion_diffusion.diffusion.matchers.LatentMatcher - - eeg_head: - enabled: True - output_dim: 14 - hidden_dim: 256 - dropout: 0.5 - pooling: mean - detach_prediction_emotion: True - use_speaker_audio: True - use_speaker_emotion: True - use_speaker_3dmm: True - use_prediction_emotion: True - - diffusion_prior: - type: - DiffusionPriorNetwork - args: - emb_preprocessing: normalize - freeze_encoder: True - audio_dim: 768 - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - _3dmm_dim: 58 - speaker_emb_dim: 512 - latent_dim: 512 - depth: 4 - num_time_layers: 2 - num_time_embeds: 1 - num_time_emb_channels: 64 - time_last_act: False - use_learned_query: True - s_audio_cond_drop_prob: 0.2 - s_latentemb_cond_drop_prob: 1.0 - s_3dmm_cond_drop_prob: 0.2 - guidance_scale: 1.0 - dim_head: 64 - heads: 8 - ff_mult: 4 - norm_in: False - norm_out: True - attn_dropout: 0.0 - ff_dropout: 0.0 - final_proj: True - normformer: False - rotary_emb: True - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - - diffusion_decoder: - type: - TransformerDenoiser - args: - emb_preprocessing: normalize - freeze_encoder: True - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - encode_emotion: False - encode_3dmm: False - ablation_skip_connection: True - nfeats: 25 - latent_dim: 512 - ff_size: 1024 - num_layers: 9 # 7 | 9 - num_heads: 8 # 4 | 8 - dropout: 0.1 - normalize_before: False - activation: gelu - flip_sin_to_cos: True - return_intermediate_dec: False - position_embedding: learned - arch: trans_dec - freq_shift: 0 - time_encoded_dim: 64 - s_audio_dim: 768 - s_audio_scale: 1.0 - s_emotion_dim: 25 - l_embed_dim: 512 - s_embed_dim: 512 - personal_emb_dim: 512 - s_3dmm_dim: 58 - concat: concat_first - guidance_scale: 1.0 # 7.5 - s_audio_enc_drop_prob: 0.2 - s_latent_embed_drop_prob: 1.0 - s_3dmm_enc_drop_prob: 0.2 - s_emotion_enc_drop_prob: 0.2 - past_l_emotion_drop_prob: 1.0 - use_past_frames: False - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - -defaults: - - /model/losses/motion_diffusion@loss - - /model/motion_diffusion/audio_embedder@audio_encoder - #- /model/latent_embedder@latent_embedder diff --git a/personalised/code/configs/model/motion_diffusion/audio_embedder.yaml b/personalised/code/configs/model/motion_diffusion/audio_embedder.yaml deleted file mode 100644 index 8e18655a404101b4ee85faa3956a14c00ea1cc4b..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/motion_diffusion/audio_embedder.yaml +++ /dev/null @@ -1,4 +0,0 @@ -type: AudioEmbedder -_target_: dataset.modules.audio_embedder.AudioEmbedder -checkpoint_path: "" -skip_norm: True \ No newline at end of file diff --git a/personalised/code/configs/model/motion_diffusion_causal.yaml b/personalised/code/configs/model/motion_diffusion_causal.yaml deleted file mode 100644 index 1317b290021f3a19e682a75d4fa53142b942a3f6..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/motion_diffusion_causal.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# Causal / coarse-to-fine variant of the generic motion_diffusion model. -# Inherits the full baseline model config and only swaps the matcher target -# and adds the causal-decoder hyperparameters. Non-invasive: originals untouched. -defaults: - - motion_diffusion - - _self_ - -diff_model: - _target_: framework.motion_diffusion.diffusion.matchers_causal.CausalLatentMatcher - diffusion_decoder: - args: - # causal lead-lag cross-attention bias - use_lag_bias: True - lag_max: 60 # max past speaker offset with its own learned bias bucket - lag_lookahead: 0 # 0 = strict causal (no future speaker frames) - # coarse-to-fine 8-class facial-expression plan - use_coarse: True - coarse_classes: 8 - coarse_hidden: 256 - coarse_emo_start: 17 # emotion 25-d layout: [0:15]=AU, [15:17]=VA, [17:25]=8 expr diff --git a/personalised/code/configs/model/motion_diffusion_velocity.yaml b/personalised/code/configs/model/motion_diffusion_velocity.yaml deleted file mode 100644 index d207d42dacd170cb509dcac0c4d5876edf9c6eda..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/motion_diffusion_velocity.yaml +++ /dev/null @@ -1,116 +0,0 @@ -diff_model: - model_name: LatentMLPMatcher - _target_: framework.motion_diffusion.diffusion.matchers_velocity.VelocityLatentMatcher - - eeg_head: - enabled: True - output_dim: 14 - hidden_dim: 256 - dropout: 0.5 - pooling: mean - detach_prediction_emotion: True - use_speaker_audio: True - use_speaker_emotion: True - use_speaker_3dmm: True - use_prediction_emotion: True - - diffusion_prior: - type: - DiffusionPriorNetwork - args: - emb_preprocessing: normalize - freeze_encoder: True - audio_dim: 768 - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - _3dmm_dim: 58 - speaker_emb_dim: 512 - latent_dim: 512 - depth: 4 - num_time_layers: 2 - num_time_embeds: 1 - num_time_emb_channels: 64 - time_last_act: False - use_learned_query: True - s_audio_cond_drop_prob: 0.2 - s_latentemb_cond_drop_prob: 1.0 - s_3dmm_cond_drop_prob: 0.2 - guidance_scale: 1.0 - dim_head: 64 - heads: 8 - ff_mult: 4 - norm_in: False - norm_out: True - attn_dropout: 0.0 - ff_dropout: 0.0 - final_proj: True - normformer: False - rotary_emb: True - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - - diffusion_decoder: - type: - TransformerDenoiser - args: - emb_preprocessing: normalize - freeze_encoder: True - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - encode_emotion: False - encode_3dmm: False - ablation_skip_connection: True - nfeats: 25 - latent_dim: 512 - ff_size: 1024 - num_layers: 9 # 7 | 9 - num_heads: 8 # 4 | 8 - dropout: 0.1 - normalize_before: False - activation: gelu - flip_sin_to_cos: True - return_intermediate_dec: False - position_embedding: learned - arch: trans_dec - freq_shift: 0 - time_encoded_dim: 64 - s_audio_dim: 768 - s_audio_scale: 1.0 - s_emotion_dim: 25 - l_embed_dim: 512 - s_embed_dim: 512 - personal_emb_dim: 512 - s_3dmm_dim: 58 - concat: concat_first - guidance_scale: 1.0 # 7.5 - s_audio_enc_drop_prob: 0.2 - s_latent_embed_drop_prob: 1.0 - s_3dmm_enc_drop_prob: 0.2 - s_emotion_enc_drop_prob: 0.2 - past_l_emotion_drop_prob: 1.0 - use_past_frames: False - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - -defaults: - - /model/losses/motion_diffusion@loss - - /model/motion_diffusion/audio_embedder@audio_encoder - #- /model/latent_embedder@latent_embedder diff --git a/personalised/code/configs/model/motion_transvae.yaml b/personalised/code/configs/model/motion_transvae.yaml deleted file mode 100644 index 08e0afc25b4cadc9af577a8faefb0761cdd8f722..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/motion_transvae.yaml +++ /dev/null @@ -1,24 +0,0 @@ -model_name: TransformerVAE -_target_: framework.motion_transvae.TransformerVAE.TransformerVAE - -img_size: 224 -audio_dim: 768 -output_3dmm_dim: 58 -output_emotion_dim: 25 -feature_dim: 128 -seq_len: ${data.train_dataset.clip_length} -task: ${task} -window_size: ${trainer.window_size} -device: cuda - -eeg_head: - enabled: True - output_dim: 14 - hidden_dim: 256 - dropout: 0.5 - pooling: mean - detach_prediction_emotion: True - use_speaker_audio: True - use_speaker_emotion: True - use_speaker_3dmm: True - use_prediction_emotion: True diff --git a/personalised/code/configs/model/optim/adam.yaml b/personalised/code/configs/model/optim/adam.yaml deleted file mode 100644 index f751f0d1d7f36e371fe68e3f45fdd5c6991c0b9d..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/optim/adam.yaml +++ /dev/null @@ -1,5 +0,0 @@ -_target_: torch.optim.Adam - -betas: [0.9, 0.98] -eps: 1e-09 -#lr: 1e-05 \ No newline at end of file diff --git a/personalised/code/configs/model/optim/adamw.yaml b/personalised/code/configs/model/optim/adamw.yaml deleted file mode 100644 index 6c06aa42da9169cc696942320b80ead3e73206a7..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/optim/adamw.yaml +++ /dev/null @@ -1,7 +0,0 @@ -_target_: torch.optim.AdamW - -#args: -#lr: 0.0001 -weight_decay: 1e-4 -betas: [0.9, 0.999] # [0.9, 0.98] -# eps: 1e-09 \ No newline at end of file diff --git a/personalised/code/configs/model/optim/sgd.yaml b/personalised/code/configs/model/optim/sgd.yaml deleted file mode 100644 index 18831c513bded4af228da868f302e6f12e165559..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/optim/sgd.yaml +++ /dev/null @@ -1,5 +0,0 @@ -_target_: torch.optim.SGD - -#lr: 0.001 -weight_decay: 1e-4 -momentum: 0.9 \ No newline at end of file diff --git a/personalised/code/configs/model/scheduler/cosine_annealing.yaml b/personalised/code/configs/model/scheduler/cosine_annealing.yaml deleted file mode 100644 index 4178cbbf8eb8d4ed7ba95d463d18f49075e6db79..0000000000000000000000000000000000000000 --- a/personalised/code/configs/model/scheduler/cosine_annealing.yaml +++ /dev/null @@ -1,4 +0,0 @@ -_target_: torch.optim.lr_scheduler.CosineAnnealingLR - -eta_min: 0 -last_epoch: -1 \ No newline at end of file diff --git a/personalised/code/configs/path.yaml b/personalised/code/configs/path.yaml deleted file mode 100644 index e510054eda62eee163e6383cd716c54b226ce784..0000000000000000000000000000000000000000 --- a/personalised/code/configs/path.yaml +++ /dev/null @@ -1,5 +0,0 @@ -# path to additional modules, defined in prepare.py -deps: ${code_path:./deps} -datasets: ${code_path:./datasets} -code_dir: ${code_path:} -working_dir: ${working_path:""} diff --git a/personalised/code/configs/pers_offline_causal.yaml b/personalised/code/configs/pers_offline_causal.yaml deleted file mode 100644 index 5ed98c9e0b85c0ef38972c73e741cceb0a603fc1..0000000000000000000000000000000000000000 --- a/personalised/code/configs/pers_offline_causal.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# Root-level top config for OFFLINE personalized (perfrdiff rewrite-weight) so it -# composes at the global package (the in-tree personalized_*/ configs get packaged -# under their subdir name when loaded via --config-name, which breaks overrides). -# Non-invasive: originals untouched; this only re-hosts the same defaults at root. -defaults: - - /personalized_offline/data/perfrdiff_rewrite_weight@data - - /personalized_offline/trainer/perfrdiff_rewrite_weight@trainer - - /shared/logger: none - - /shared/path@path - - _self_ - -seed: 1234 -logger_level: INFO -task: offline -stage: ??? -data_dir: ??? -resume: false -run_id: ${generate_id:} -resume_id: "" -num_gts: 10 diff --git a/personalised/code/configs/pers_online_causal.yaml b/personalised/code/configs/pers_online_causal.yaml deleted file mode 100644 index 2f2e3e37417144295566cd16932153b5ca605bcb..0000000000000000000000000000000000000000 --- a/personalised/code/configs/pers_online_causal.yaml +++ /dev/null @@ -1,19 +0,0 @@ -# Root-level top config for ONLINE personalized (perfrdiff rewrite-weight), mirror of -# pers_offline_causal.yaml: re-hosts the personalized_online defaults at the global -# package so --config-name overrides compose correctly. Non-invasive. -defaults: - - /personalized_online/data/perfrdiff_rewrite_weight@data - - /personalized_online/trainer/perfrdiff_rewrite_weight@trainer - - /shared/logger: none - - /shared/path@path - - _self_ - -seed: 1234 -logger_level: INFO -task: online -stage: ??? -data_dir: ??? -resume: false -run_id: ${generate_id:} -resume_id: "" -num_gts: 10 diff --git a/personalised/code/configs/personalized_offline/data/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/personalized_offline/data/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index 06e8e5fe9131535e98886a7e3f694dc46023a1c8..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_offline/data/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,135 +0,0 @@ -data_name: react_2025 -_target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataModule - -train_dataset: - _target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 8 - root_dir: ${data_dir} - split: train - clip_length: 750 - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_ref: True - load_personality_l: True - personal_condition_mode: ${trainer.main_model.args.personal_condition_mode} - load_eeg_l: ${trainer.generic.train_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - normalize_3dmm: standard - -validation_dataset: - _target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 8 - root_dir: ${data_dir} - split: val - clip_length: 750 - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_ref: True - load_personality_l: True - personal_condition_mode: ${trainer.main_model.args.personal_condition_mode} - load_eeg_l: ${trainer.generic.train_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - normalize_3dmm: standard - -test_dataset: - _target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 4 - root_dir: ${data_dir} - split: test - clip_length: 750 - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_ref: False - load_personality_l: True - personal_condition_mode: ${trainer.main_model.args.personal_condition_mode} - load_eeg_l: ${trainer.generic.eval_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - normalize_3dmm: standard - num_test_gts: ${num_gts} diff --git a/personalised/code/configs/personalized_offline/model/losses/motion_diffusion.yaml b/personalised/code/configs/personalized_offline/model/losses/motion_diffusion.yaml deleted file mode 100644 index f528b1e1cb4be34a38fde02543cc3945f2dd520f..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_offline/model/losses/motion_diffusion.yaml +++ /dev/null @@ -1,11 +0,0 @@ -_target_: framework.utils.losses.DiffusionLoss - -losses_type: MSELossApt -losses_multiplier: 1.0 -n_preds: ${trainer.${trainer.trainer_mode}.num_preds} -temporal_loss_w: 0.0 -prior_loss_weight: 1.0 -eeg_loss_weight: 0.25 -w_au: 1.0 # action unit -w_va: 5.0 # valence and arousal -w_em: 5.0 # emotion diff --git a/personalised/code/configs/personalized_offline/model/losses/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/personalized_offline/model/losses/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index a8cd2da396773e073c381619488d0031e429476a..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_offline/model/losses/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,8 +0,0 @@ -type: DiffusionLoss -args: - losses_type: [MSELoss, MSELoss] - losses_multipliers: [0, 1] - losses_decoded: [False, True] - k: ${trainer.${trainer.trainer_mode}.num_preds} - temporal_loss_w: 0.0 - eeg_loss_weight: 1.0 diff --git a/personalised/code/configs/personalized_offline/model/motion_diffusion.yaml b/personalised/code/configs/personalized_offline/model/motion_diffusion.yaml deleted file mode 100644 index 7dcde62f171fefd7caa03974090625e9af068ff2..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_offline/model/motion_diffusion.yaml +++ /dev/null @@ -1,116 +0,0 @@ -diff_model: - model_name: LatentMLPMatcher - _target_: framework.motion_diffusion.diffusion.matchers.LatentMatcher - - eeg_head: - enabled: True - output_dim: 14 - hidden_dim: 256 - dropout: 0.5 - pooling: mean - detach_prediction_emotion: True - use_speaker_audio: True - use_speaker_emotion: True - use_speaker_3dmm: True - use_prediction_emotion: True - - diffusion_prior: - type: - DiffusionPriorNetwork - args: - emb_preprocessing: normalize - freeze_encoder: True - audio_dim: 768 - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - _3dmm_dim: 58 - speaker_emb_dim: 512 - latent_dim: 512 - depth: 4 - num_time_layers: 2 - num_time_embeds: 1 - num_time_emb_channels: 64 - time_last_act: False - use_learned_query: True - s_audio_cond_drop_prob: 0.2 - s_latentemb_cond_drop_prob: 1.0 - s_3dmm_cond_drop_prob: 0.2 - guidance_scale: 1.0 - dim_head: 64 - heads: 8 - ff_mult: 4 - norm_in: False - norm_out: True - attn_dropout: 0.0 - ff_dropout: 0.0 - final_proj: True - normformer: False - rotary_emb: True - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - - diffusion_decoder: - type: - TransformerDenoiser - args: - emb_preprocessing: normalize - freeze_encoder: True - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - encode_emotion: False - encode_3dmm: False - ablation_skip_connection: True - nfeats: 25 - latent_dim: 512 - ff_size: 1024 - num_layers: 9 # 7 | 9 - num_heads: 8 # 4 | 8 - dropout: 0.1 - normalize_before: False - activation: gelu - flip_sin_to_cos: True - return_intermediate_dec: False - position_embedding: learned - arch: trans_dec - freq_shift: 0 - time_encoded_dim: 64 - s_audio_dim: 768 - s_audio_scale: 1.0 - s_emotion_dim: 25 - l_embed_dim: 512 - s_embed_dim: 512 - personal_emb_dim: 512 - s_3dmm_dim: 58 - concat: concat_first - guidance_scale: 1.0 # 7.5 - s_audio_enc_drop_prob: 0.2 - s_latent_embed_drop_prob: 1.0 - s_3dmm_enc_drop_prob: 0.2 - s_emotion_enc_drop_prob: 0.2 - past_l_emotion_drop_prob: 1.0 - use_past_frames: False - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - -defaults: - - /personalized_offline/model/losses/motion_diffusion@loss - - /personalized_offline/model/motion_diffusion/audio_embedder@audio_encoder - #- /model/latent_embedder@latent_embedder diff --git a/personalised/code/configs/personalized_offline/model/motion_diffusion/audio_embedder.yaml b/personalised/code/configs/personalized_offline/model/motion_diffusion/audio_embedder.yaml deleted file mode 100644 index 8e18655a404101b4ee85faa3956a14c00ea1cc4b..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_offline/model/motion_diffusion/audio_embedder.yaml +++ /dev/null @@ -1,4 +0,0 @@ -type: AudioEmbedder -_target_: dataset.modules.audio_embedder.AudioEmbedder -checkpoint_path: "" -skip_norm: True \ No newline at end of file diff --git a/personalised/code/configs/personalized_offline/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/personalized_offline/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index d4c9ab4249c5046bfcd830950389218e3d278aac..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_offline/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,23 +0,0 @@ -hydra: - run: - dir: outputs/${trainer.model_name}/${data.data_name}/${task}/${run_id} - job: - chdir: true - -seed: 1234 -logger_level: INFO - -task: offline -stage: ??? -data_dir: ??? -resume: false -run_id: ${generate_id:} -resume_id: "" -num_gts: 10 - -defaults: - - /personalized_offline/data/perfrdiff_rewrite_weight@data - - /personalized_offline/trainer/perfrdiff_rewrite_weight@trainer - - /shared/logger: none - - /shared/path@path - - _self_ diff --git a/personalised/code/configs/personalized_offline/trainer/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/personalized_offline/trainer/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index a7d28d15d69518a5fa44b17e0ae4695c87817bf8..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_offline/trainer/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,89 +0,0 @@ -model_name: perfrdiff_rewrite_weight -_target_: trainer.perfrdiff_rewrite_weight.Trainer - -trainer_mode: generic -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} -current_runid: ${run_id} -resume_runid: ${resume_id} -batch_size: 4 -post_config_name: configs/shared/model/emotion_autoencoder.yaml -post_clip_length: 1000 -data_clamp: True -num_eval_preds: 10 -eval_clip_batch_size: 16 -parallel_eval_preds: True - -pretrained: - diffusion_prior: pretrained_models/diffusion_model/DiffusionPriorNetwork/checkpoint.pth - diffusion_decoder: pretrained_models/diffusion_model/TransformerDenoiser/checkpoint.pth - modifier_checkpoint: "" - eeg_head_checkpoint: pretrained_models/rewrite_weight/EEGPredictionHead/checkpoint.pth - -person_specific: - checkpoint_path: pretrained_models/person_specific/checkpoint.pth - args: - in_features: 58 - embed_dim: 512 - num_heads: 4 - num_layers: 4 - mlp_dim: 1024 - seq_len: 750 - proj_dim: 512 - proj_head: mlp - drop_prob: 0.1 - max_len: 1000 - pos_encoding: absolute - embed_layer: linear - -main_model: - args: - input_dim: 512 - latent_dim: 1024 - embed_dim: 512 - personal_condition_mode: 3dmm_personality - personality_input_dim: 5 - personality_hidden_dim: 128 - personality_dropout: 0.1 - personality_fusion_hidden_dim: 512 - regularization: False - regular_w: 0.0 - num_shared_layers: 2 - modified_layers: - - diffusion_decoder.model.decoder.layers.4.multihead_attn - - diffusion_decoder.model.decoder.layers.6.multihead_attn - - diffusion_decoder.model.to_emotion_feat - predict: shift - modify: all - optimizer_hypernet: - type: sgd - args: - lr: 0.001 - weight_decay: 1e-4 - momentum: 0.9 - -generic: - seed: 1234 - start_epoch: 0 - epochs: 100 - model: LatentMatcher - clip_grad: False - num_workers: 16 - log_dir: log - tb_dir: tb_logs - out_dir: results - save_period: 10 - val_period: 10 - num_preds: 1 - clip_length: 750 - window_size: 30 - s_ratio: 2 - train_eeg: False - train_eeg_head_only: False - eval_eeg: True - bidirectional: False - -defaults: - - /personalized_offline/model/motion_diffusion@model - - /personalized_offline/model/losses/perfrdiff_rewrite_weight@criterion diff --git a/personalised/code/configs/personalized_online/data/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/personalized_online/data/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index 06e8e5fe9131535e98886a7e3f694dc46023a1c8..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_online/data/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,135 +0,0 @@ -data_name: react_2025 -_target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataModule - -train_dataset: - _target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 8 - root_dir: ${data_dir} - split: train - clip_length: 750 - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_ref: True - load_personality_l: True - personal_condition_mode: ${trainer.main_model.args.personal_condition_mode} - load_eeg_l: ${trainer.generic.train_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - normalize_3dmm: standard - -validation_dataset: - _target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 8 - root_dir: ${data_dir} - split: val - clip_length: 750 - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_ref: True - load_personality_l: True - personal_condition_mode: ${trainer.main_model.args.personal_condition_mode} - load_eeg_l: ${trainer.generic.train_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - normalize_3dmm: standard - -test_dataset: - _target_: dataset.perfrdiff_rewrite_weight.PerFRDiffRewriteWeightDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 4 - root_dir: ${data_dir} - split: test - clip_length: 750 - target_size: 224 - crop_size: 224 - fps: 30 - audio_feature_type: wav2vec - load_video_s: False - load_video_l: False - load_audio: True - load_emotion_s: True - load_emotion_l: True - load_3dmm_s: True - load_3dmm_l: True - load_ref: False - load_personality_l: True - personal_condition_mode: ${trainer.main_model.args.personal_condition_mode} - load_eeg_l: ${trainer.generic.eval_eeg} - eeg_dir_name: eeg_processed - eeg_target_cols: - - TP9 - - AF7 - - AF8 - - TP10 - - Delta_TP9 - - Theta_TP9 - - Alpha_TP9 - - Beta_TP9 - - Gamma_TP9 - - Delta_TP10 - - Theta_TP10 - - Alpha_TP10 - - Beta_TP10 - - Gamma_TP10 - eeg_channel_scale: 1000.0 - eeg_use_tar_fallback: True - bidirectional: ${trainer.generic.bidirectional} - normalize_3dmm: standard - num_test_gts: ${num_gts} diff --git a/personalised/code/configs/personalized_online/model/losses/motion_diffusion.yaml b/personalised/code/configs/personalized_online/model/losses/motion_diffusion.yaml deleted file mode 100644 index f528b1e1cb4be34a38fde02543cc3945f2dd520f..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_online/model/losses/motion_diffusion.yaml +++ /dev/null @@ -1,11 +0,0 @@ -_target_: framework.utils.losses.DiffusionLoss - -losses_type: MSELossApt -losses_multiplier: 1.0 -n_preds: ${trainer.${trainer.trainer_mode}.num_preds} -temporal_loss_w: 0.0 -prior_loss_weight: 1.0 -eeg_loss_weight: 0.25 -w_au: 1.0 # action unit -w_va: 5.0 # valence and arousal -w_em: 5.0 # emotion diff --git a/personalised/code/configs/personalized_online/model/losses/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/personalized_online/model/losses/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index a8cd2da396773e073c381619488d0031e429476a..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_online/model/losses/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,8 +0,0 @@ -type: DiffusionLoss -args: - losses_type: [MSELoss, MSELoss] - losses_multipliers: [0, 1] - losses_decoded: [False, True] - k: ${trainer.${trainer.trainer_mode}.num_preds} - temporal_loss_w: 0.0 - eeg_loss_weight: 1.0 diff --git a/personalised/code/configs/personalized_online/model/motion_diffusion.yaml b/personalised/code/configs/personalized_online/model/motion_diffusion.yaml deleted file mode 100644 index 3bddb28da79a5570ca657a268db3e9b4f484b311..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_online/model/motion_diffusion.yaml +++ /dev/null @@ -1,116 +0,0 @@ -diff_model: - model_name: LatentMLPMatcher - _target_: framework.motion_diffusion.diffusion.matchers.LatentMatcher - - eeg_head: - enabled: True - output_dim: 14 - hidden_dim: 256 - dropout: 0.5 - pooling: mean - detach_prediction_emotion: True - use_speaker_audio: True - use_speaker_emotion: True - use_speaker_3dmm: True - use_prediction_emotion: True - - diffusion_prior: - type: - DiffusionPriorNetwork - args: - emb_preprocessing: normalize - freeze_encoder: True - audio_dim: 768 - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - _3dmm_dim: 58 - speaker_emb_dim: 512 - latent_dim: 512 - depth: 4 - num_time_layers: 2 - num_time_embeds: 1 - num_time_emb_channels: 64 - time_last_act: False - use_learned_query: True - s_audio_cond_drop_prob: 0.2 - s_latentemb_cond_drop_prob: 1.0 - s_3dmm_cond_drop_prob: 0.2 - guidance_scale: 1.0 - dim_head: 64 - heads: 8 - ff_mult: 4 - norm_in: False - norm_out: True - attn_dropout: 0.0 - ff_dropout: 0.0 - final_proj: True - normformer: False - rotary_emb: True - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - - diffusion_decoder: - type: - TransformerDenoiser - args: - emb_preprocessing: normalize - freeze_encoder: True - window_size: ${trainer.${trainer.trainer_mode}.window_size} - s_ratio: ${trainer.${trainer.trainer_mode}.s_ratio} - token_len: 750 - encode_emotion: False - encode_3dmm: False - ablation_skip_connection: True - nfeats: 25 - latent_dim: 512 - ff_size: 1024 - num_layers: 9 # 7 | 9 - num_heads: 8 # 4 | 8 - dropout: 0.1 - normalize_before: False - activation: gelu - flip_sin_to_cos: True - return_intermediate_dec: False - position_embedding: learned - arch: trans_dec - freq_shift: 0 - time_encoded_dim: 64 - s_audio_dim: 768 - s_audio_scale: 1.0 - s_emotion_dim: 25 - l_embed_dim: 512 - s_embed_dim: 512 - personal_emb_dim: 512 - s_3dmm_dim: 58 - concat: concat_first - guidance_scale: 1.0 # 7.5 - s_audio_enc_drop_prob: 0.2 - s_latent_embed_drop_prob: 1.0 - s_3dmm_enc_drop_prob: 0.2 - s_emotion_enc_drop_prob: 0.2 - past_l_emotion_drop_prob: 1.0 - use_past_frames: False - scheduler: - noise_schedule: cosine - timestep_spacing: leading - num_train_timesteps: 1000 - num_inference_timesteps: 50 - predict: start_x - var_type: fixed_large - rescale_timesteps: False - noise_std: 1 - num_preds: ${trainer.${trainer.trainer_mode}.num_preds} - -defaults: - - /personalized_online/model/losses/motion_diffusion@loss - - /personalized_online/model/motion_diffusion/audio_embedder@audio_encoder - #- /model/latent_embedder@latent_embedder diff --git a/personalised/code/configs/personalized_online/model/motion_diffusion/audio_embedder.yaml b/personalised/code/configs/personalized_online/model/motion_diffusion/audio_embedder.yaml deleted file mode 100644 index 8e18655a404101b4ee85faa3956a14c00ea1cc4b..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_online/model/motion_diffusion/audio_embedder.yaml +++ /dev/null @@ -1,4 +0,0 @@ -type: AudioEmbedder -_target_: dataset.modules.audio_embedder.AudioEmbedder -checkpoint_path: "" -skip_norm: True \ No newline at end of file diff --git a/personalised/code/configs/personalized_online/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/personalized_online/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index c50bacd98086a6d0354ba1b620da7f7805de182a..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_online/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,23 +0,0 @@ -hydra: - run: - dir: outputs/${trainer.model_name}/${data.data_name}/${task}/${run_id} - job: - chdir: true - -seed: 1234 -logger_level: INFO - -task: online -stage: ??? -data_dir: ??? -resume: false -run_id: ${generate_id:} -resume_id: "" -num_gts: 10 - -defaults: - - /personalized_online/data/perfrdiff_rewrite_weight@data - - /personalized_online/trainer/perfrdiff_rewrite_weight@trainer - - /shared/logger: none - - /shared/path@path - - _self_ diff --git a/personalised/code/configs/personalized_online/trainer/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/personalized_online/trainer/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index 0129543e100797b74c0cf6c2616f772c527c78bc..0000000000000000000000000000000000000000 --- a/personalised/code/configs/personalized_online/trainer/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,89 +0,0 @@ -model_name: perfrdiff_rewrite_weight -_target_: trainer.perfrdiff_rewrite_weight.Trainer - -trainer_mode: generic -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} -current_runid: ${run_id} -resume_runid: ${resume_id} -batch_size: 4 -post_config_name: configs/shared/model/emotion_autoencoder.yaml -post_clip_length: 1000 -data_clamp: True -num_eval_preds: 10 -eval_clip_batch_size: 16 -parallel_eval_preds: True - -pretrained: - diffusion_prior: pretrained_models/diffusion_model/DiffusionPriorNetwork/checkpoint.pth - diffusion_decoder: pretrained_models/diffusion_model/TransformerDenoiser/checkpoint.pth - modifier_checkpoint: "" - eeg_head_checkpoint: pretrained_models/rewrite_weight/EEGPredictionHead/checkpoint.pth - -person_specific: - checkpoint_path: pretrained_models/person_specific/checkpoint.pth - args: - in_features: 58 - embed_dim: 512 - num_heads: 4 - num_layers: 4 - mlp_dim: 1024 - seq_len: 750 - proj_dim: 512 - proj_head: mlp - drop_prob: 0.1 - max_len: 1000 - pos_encoding: absolute - embed_layer: linear - -main_model: - args: - input_dim: 512 - latent_dim: 1024 - embed_dim: 512 - personal_condition_mode: 3dmm_personality - personality_input_dim: 5 - personality_hidden_dim: 128 - personality_dropout: 0.1 - personality_fusion_hidden_dim: 512 - regularization: False - regular_w: 0.0 - num_shared_layers: 2 - modified_layers: - - diffusion_decoder.model.decoder.layers.4.multihead_attn - - diffusion_decoder.model.decoder.layers.6.multihead_attn - - diffusion_decoder.model.to_emotion_feat - predict: shift - modify: all - optimizer_hypernet: - type: sgd - args: - lr: 0.001 - weight_decay: 1e-4 - momentum: 0.9 - -generic: - seed: 1234 - start_epoch: 0 - epochs: 100 - model: LatentMatcher - clip_grad: False - num_workers: 16 - log_dir: log - tb_dir: tb_logs - out_dir: results - save_period: 10 - val_period: 10 - num_preds: 1 - clip_length: 750 - window_size: 30 - s_ratio: 2 - train_eeg: False - train_eeg_head_only: False - eval_eeg: True - bidirectional: False - -defaults: - - /personalized_online/model/motion_diffusion@model - - /personalized_online/model/losses/perfrdiff_rewrite_weight@criterion diff --git a/personalised/code/configs/shared/data/emotion_autoencoder.yaml b/personalised/code/configs/shared/data/emotion_autoencoder.yaml deleted file mode 100644 index c1578cc203b43cf4f693650cea47a429b13937df..0000000000000000000000000000000000000000 --- a/personalised/code/configs/shared/data/emotion_autoencoder.yaml +++ /dev/null @@ -1,29 +0,0 @@ -data_name: react_2025 -_target_: dataset.react_2025.ReactionDataloader - -train_dataset: - _target_: dataset.react_2025.ReactionAutoEncoderDataset - batch_size: ${trainer.batch_size} - shuffle: True - num_workers: 2 - root_dir: ${data_dir} - split: train - clip_length: ${trainer.clip_length} - -validation_dataset: - _target_: dataset.react_2025.ReactionAutoEncoderDataset - batch_size: ${trainer.batch_size} - shuffle: False - num_workers: 2 - root_dir: ${data_dir} - split: val - clip_length: ${trainer.clip_length} - -test_dataset: - _target_: dataset.react_2025.ReactionAutoEncoderDataset - batch_size: 1 - shuffle: False - num_workers: 2 - root_dir: ${data_dir} - split: test - clip_length: ${trainer.clip_length} \ No newline at end of file diff --git a/personalised/code/configs/shared/logger/none.yaml b/personalised/code/configs/shared/logger/none.yaml deleted file mode 100644 index 3da2aede1c03e6d9660ccabe8ce187288aee3f48..0000000000000000000000000000000000000000 --- a/personalised/code/configs/shared/logger/none.yaml +++ /dev/null @@ -1,4 +0,0 @@ -logger_name: none -version: ${run_id} -logger_dir: tb_logs -project: null \ No newline at end of file diff --git a/personalised/code/configs/shared/model/emotion_autoencoder.yaml b/personalised/code/configs/shared/model/emotion_autoencoder.yaml deleted file mode 100644 index 2f80bbb7189d799d036a01b075ba3cc332a83b03..0000000000000000000000000000000000000000 --- a/personalised/code/configs/shared/model/emotion_autoencoder.yaml +++ /dev/null @@ -1,17 +0,0 @@ -type: EmotionVAE -_target_: framework.modules.emotion_autoencoder.EmotionVAE - -in_channels: 25 -out_channels: 25 -feature_dim: 512 -nhead: 8 -dropout: 0.1 -num_encoder_layers: 6 -num_decoder_layers: 6 -mlp_dist: False # expand mu & logvar -in_proj_type : linear # linear | mlp -out_proj_type : separate # separate | shared -pe_type : absolute # learnable | absolute -query_type : zero -max_seq_len: 5000 -global_token_len: 64 \ No newline at end of file diff --git a/personalised/code/configs/shared/model/losses/emotion_autoencoder.yaml b/personalised/code/configs/shared/model/losses/emotion_autoencoder.yaml deleted file mode 100644 index 6881880215136f33e15ae18141d3470b69d6afc2..0000000000000000000000000000000000000000 --- a/personalised/code/configs/shared/model/losses/emotion_autoencoder.yaml +++ /dev/null @@ -1,6 +0,0 @@ -_target_: framework.utils.losses.EmotionVAELoss - -w_au: 1.0 -w_va: 10.0 -w_em: 1.0 -w_kld: 0.001 \ No newline at end of file diff --git a/personalised/code/configs/shared/model/optim/adam.yaml b/personalised/code/configs/shared/model/optim/adam.yaml deleted file mode 100644 index f751f0d1d7f36e371fe68e3f45fdd5c6991c0b9d..0000000000000000000000000000000000000000 --- a/personalised/code/configs/shared/model/optim/adam.yaml +++ /dev/null @@ -1,5 +0,0 @@ -_target_: torch.optim.Adam - -betas: [0.9, 0.98] -eps: 1e-09 -#lr: 1e-05 \ No newline at end of file diff --git a/personalised/code/configs/shared/model/optim/adamw.yaml b/personalised/code/configs/shared/model/optim/adamw.yaml deleted file mode 100644 index 6c06aa42da9169cc696942320b80ead3e73206a7..0000000000000000000000000000000000000000 --- a/personalised/code/configs/shared/model/optim/adamw.yaml +++ /dev/null @@ -1,7 +0,0 @@ -_target_: torch.optim.AdamW - -#args: -#lr: 0.0001 -weight_decay: 1e-4 -betas: [0.9, 0.999] # [0.9, 0.98] -# eps: 1e-09 \ No newline at end of file diff --git a/personalised/code/configs/shared/model/optim/sgd.yaml b/personalised/code/configs/shared/model/optim/sgd.yaml deleted file mode 100644 index 18831c513bded4af228da868f302e6f12e165559..0000000000000000000000000000000000000000 --- a/personalised/code/configs/shared/model/optim/sgd.yaml +++ /dev/null @@ -1,5 +0,0 @@ -_target_: torch.optim.SGD - -#lr: 0.001 -weight_decay: 1e-4 -momentum: 0.9 \ No newline at end of file diff --git a/personalised/code/configs/shared/model/scheduler/cosine_annealing.yaml b/personalised/code/configs/shared/model/scheduler/cosine_annealing.yaml deleted file mode 100644 index 4178cbbf8eb8d4ed7ba95d463d18f49075e6db79..0000000000000000000000000000000000000000 --- a/personalised/code/configs/shared/model/scheduler/cosine_annealing.yaml +++ /dev/null @@ -1,4 +0,0 @@ -_target_: torch.optim.lr_scheduler.CosineAnnealingLR - -eta_min: 0 -last_epoch: -1 \ No newline at end of file diff --git a/personalised/code/configs/shared/path.yaml b/personalised/code/configs/shared/path.yaml deleted file mode 100644 index e510054eda62eee163e6383cd716c54b226ce784..0000000000000000000000000000000000000000 --- a/personalised/code/configs/shared/path.yaml +++ /dev/null @@ -1,5 +0,0 @@ -# path to additional modules, defined in prepare.py -deps: ${code_path:./deps} -datasets: ${code_path:./datasets} -code_dir: ${code_path:} -working_dir: ${working_path:""} diff --git a/personalised/code/configs/shared/trainer/emotion_autoencoder.yaml b/personalised/code/configs/shared/trainer/emotion_autoencoder.yaml deleted file mode 100644 index 33721984a93040d270b2a96608b75ea1db4ed9b4..0000000000000000000000000000000000000000 --- a/personalised/code/configs/shared/trainer/emotion_autoencoder.yaml +++ /dev/null @@ -1,23 +0,0 @@ -model_name: emotion_autoencoder -_target_: trainer.emotion_autoencoder.Trainer - -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} -current_runid: ${run_id} -resume_runid: ${resume_id} - -start_epoch: 0 -epochs: 3000 -tb_dir: tb_log -val_period: 50 -clip_length: 1000 -batch_size: 8 -lr: 0.0001 -weight_decay: 1e-4 -beta: [0.9, 0.999] - -defaults: - - /shared/model/emotion_autoencoder@model - - /shared/model/optim/adamw@optim - - /shared/model/losses/emotion_autoencoder@criterion # loss func diff --git a/personalised/code/configs/trainer/_archive.yaml b/personalised/code/configs/trainer/_archive.yaml deleted file mode 100644 index 703c7e0c9966ae65665657a43c33cf98360e8ee1..0000000000000000000000000000000000000000 --- a/personalised/code/configs/trainer/_archive.yaml +++ /dev/null @@ -1,105 +0,0 @@ -## Trainer Configs -#model_name: latent_embedder -#_target_: trainer.latent_embedder.Trainer -#folder: save/${trainer.model_name}/${data.data_name} -#ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -#resumed_training: ${resume} -#current_runid: ${run_id} -#resume_runid: ${resume_id} -#start_epoch: 0 -#epochs: 200 -#val_period: 10 -#lr: 0.001 -#weight_decay: 5e-4 -#beta: [0.9, 0.999] -#clip_length: 750 -#window_size: 30 -#s_ratio: 2 -#tb_dir: tb_logs -#batch_size: 16 -#defaults: -# - /model/latent_embedder@model -# - /model/optim/adamw@optim -# - /model/losses/latent_embedder@criterion # loss func - -## Model Configs -#type: LatentEmbedder -#_target_: framework.motion_diffusion.diffusion.rnn.LatentEmbedder -#checkpoint_path: pretrained_models/latent_embedder/checkpoint.pth -#emotion_dim: 25 -#coeff_3dmm_dim: 58 -#emb_dims: [128, 128] -#num_layers: 2 -#hidden_dim: 512 -#z_dim: 512 -#rnn_type: 'gru' -#dropout: 0.0 - -## Loss Configs -#_target_: framework.motion_diffusion.diffusion.utils.losses.MSELoss_AE -#w_mse: 1.0 -#w_kld: 0.00001 -#w_coeff: 1.0 - -## Data Configs -#data_name: react_2025 -#_target_: dataset.react_2025.ReactionDataloader -#train_dataset: -# _target_: dataset.react_2025.ReactionDataset -# batch_size: ${trainer.batch_size} -# shuffle: True -# num_workers: 16 -# root_dir: ${data_dir} -# split: train -# clip_length: 750 -# target_size: 224 -# crop_size: 224 -# fps: 30 -# audio_feature_type: wav2vec -# load_video_s: False -# load_video_l: False -# load_audio: False -# load_emotion_s: True -# load_emotion_l: False -# load_3dmm_s: True -# load_3dmm_l: False - -#validation_dataset: -# _target_: dataset.react_2025.ReactionDataset -# batch_size: ${trainer.batch_size} -# shuffle: False -# num_workers: 16 -# root_dir: ${data_dir} -# split: val -# target_size: 224 -# crop_size: 224 -# clip_length: 750 -# fps: 30 -# audio_feature_type: wav2vec -# load_video_s: False -# load_video_l: False -# load_audio: False -# load_emotion_s: True -# load_emotion_l: False -# load_3dmm_s: True -# load_3dmm_l: False - -#test_dataset: -# _target_: dataset.react_2025.ReactionDataset -# batch_size: 1 -# shuffle: False -# num_workers: 8 -# root_dir: ${data_dir} -# split: test -# target_size: 224 -# crop_size: 224 -# clip_length: 750 -# fps: 30 -# audio_feature_type: wav2vec -# load_video_s: False -# load_video_l: False -# load_audio: False -# load_emotion_s: True -# load_emotion_l: False -# load_3dmm_s: True -# load_3dmm_l: False \ No newline at end of file diff --git a/personalised/code/configs/trainer/emotion_autoencoder.yaml b/personalised/code/configs/trainer/emotion_autoencoder.yaml deleted file mode 100644 index 4df0d9bfb6ea69c75f8b0830827743dac16c0759..0000000000000000000000000000000000000000 --- a/personalised/code/configs/trainer/emotion_autoencoder.yaml +++ /dev/null @@ -1,23 +0,0 @@ -model_name: emotion_autoencoder -_target_: trainer.emotion_autoencoder.Trainer - -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} -current_runid: ${run_id} -resume_runid: ${resume_id} - -start_epoch: 0 -epochs: 3000 -tb_dir: tb_log -val_period: 50 -clip_length: 1000 -batch_size: 8 -lr: 0.0001 -weight_decay: 1e-4 -beta: [0.9, 0.999] - -defaults: - - /model/emotion_autoencoder@model - - /model/optim/adamw@optim - - /model/losses/emotion_autoencoder@criterion # loss func diff --git a/personalised/code/configs/trainer/motion_diffusion.yaml b/personalised/code/configs/trainer/motion_diffusion.yaml deleted file mode 100644 index b5ce777f5f7b9ecedbfc484ad7de0be352f25c37..0000000000000000000000000000000000000000 --- a/personalised/code/configs/trainer/motion_diffusion.yaml +++ /dev/null @@ -1,48 +0,0 @@ -model_name: motion_diffusion -_target_: trainer.motion_diffusion.Trainer - -trainer_mode: generic # generic facial reaction generator -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} # true | false -current_runid: ${run_id} -resume_runid: ${resume_id} -batch_size: 4 -post_config_name: configs/model/emotion_autoencoder.yaml -post_clip_length: 1000 # ground-truths alignment -data_clamp: True - -generic: - seed: 1234 - start_epoch: 0 - epochs: 500 - lr: 0.0001 - model: LatentMatcher - clip_grad: False - num_workers: 16 - log_dir: log - tb_dir: tb_logs - out_dir: results - save_period: 10 - val_period: 10 - num_preds: 10 - clip_length: 750 - window_size: 30 - s_ratio: 2 - train_eeg_head_only: False - pretrained_decoder_checkpoint: "" - pretrained_prior_checkpoint: "" - pretrained_load_prior: True - eval_eeg: True - bidirectional: False - -renderer: # face animation based on predicted 3DMM coefficients - name: PIRenderer - _target_: utils.render.Render - do_render: False - -defaults: - - /model/motion_diffusion@model - - /model/optim/adamw@optim - - /model/scheduler/cosine_annealing@scheduler - - /model/losses/motion_diffusion@criterion # loss func diff --git a/personalised/code/configs/trainer/motion_diffusion_causal.yaml b/personalised/code/configs/trainer/motion_diffusion_causal.yaml deleted file mode 100644 index 8380d6000738d4a50155cf6712a4ec7a0093017a..0000000000000000000000000000000000000000 --- a/personalised/code/configs/trainer/motion_diffusion_causal.yaml +++ /dev/null @@ -1,48 +0,0 @@ -model_name: motion_diffusion -_target_: trainer.motion_diffusion.Trainer - -trainer_mode: generic # generic facial reaction generator -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} # true | false -current_runid: ${run_id} -resume_runid: ${resume_id} -batch_size: 4 -post_config_name: configs/model/emotion_autoencoder.yaml -post_clip_length: 1000 # ground-truths alignment -data_clamp: True - -generic: - seed: 1234 - start_epoch: 0 - epochs: 500 - lr: 0.0001 - model: LatentMatcher - clip_grad: False - num_workers: 16 - log_dir: log - tb_dir: tb_logs - out_dir: results - save_period: 10 - val_period: 10 - num_preds: 10 - clip_length: 750 - window_size: 30 - s_ratio: 2 - train_eeg_head_only: False - pretrained_decoder_checkpoint: "" - pretrained_prior_checkpoint: "" - pretrained_load_prior: True - eval_eeg: True - bidirectional: False - -renderer: # face animation based on predicted 3DMM coefficients - name: PIRenderer - _target_: utils.render.Render - do_render: False - -defaults: - - /model/motion_diffusion_causal@model - - /model/optim/adamw@optim - - /model/scheduler/cosine_annealing@scheduler - - /model/losses/motion_diffusion_causal@criterion # loss func diff --git a/personalised/code/configs/trainer/motion_diffusion_causal_dp.yaml b/personalised/code/configs/trainer/motion_diffusion_causal_dp.yaml deleted file mode 100644 index 7acbdc23912ffcd6421e6dd76f45b5b9a60abc53..0000000000000000000000000000000000000000 --- a/personalised/code/configs/trainer/motion_diffusion_causal_dp.yaml +++ /dev/null @@ -1,48 +0,0 @@ -model_name: motion_diffusion -_target_: trainer.motion_diffusion_dp.Trainer - -trainer_mode: generic # generic facial reaction generator -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} # true | false -current_runid: ${run_id} -resume_runid: ${resume_id} -batch_size: 4 -post_config_name: configs/model/emotion_autoencoder.yaml -post_clip_length: 1000 # ground-truths alignment -data_clamp: True - -generic: - seed: 1234 - start_epoch: 0 - epochs: 500 - lr: 0.0001 - model: LatentMatcher - clip_grad: False - num_workers: 16 - log_dir: log - tb_dir: tb_logs - out_dir: results - save_period: 10 - val_period: 10 - num_preds: 10 - clip_length: 750 - window_size: 30 - s_ratio: 2 - train_eeg_head_only: False - pretrained_decoder_checkpoint: "" - pretrained_prior_checkpoint: "" - pretrained_load_prior: True - eval_eeg: True - bidirectional: False - -renderer: # face animation based on predicted 3DMM coefficients - name: PIRenderer - _target_: utils.render.Render - do_render: False - -defaults: - - /model/motion_diffusion_causal@model - - /model/optim/adamw@optim - - /model/scheduler/cosine_annealing@scheduler - - /model/losses/motion_diffusion_causal@criterion # loss func diff --git a/personalised/code/configs/trainer/motion_diffusion_dp.yaml b/personalised/code/configs/trainer/motion_diffusion_dp.yaml deleted file mode 100644 index 30348d360434e79b4f261ee4a262448916b17b2b..0000000000000000000000000000000000000000 --- a/personalised/code/configs/trainer/motion_diffusion_dp.yaml +++ /dev/null @@ -1,48 +0,0 @@ -model_name: motion_diffusion -_target_: trainer.motion_diffusion_dp.Trainer - -trainer_mode: generic # generic facial reaction generator -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} # true | false -current_runid: ${run_id} -resume_runid: ${resume_id} -batch_size: 4 -post_config_name: configs/model/emotion_autoencoder.yaml -post_clip_length: 1000 # ground-truths alignment -data_clamp: True - -generic: - seed: 1234 - start_epoch: 0 - epochs: 500 - lr: 0.0001 - model: LatentMatcher - clip_grad: False - num_workers: 16 - log_dir: log - tb_dir: tb_logs - out_dir: results - save_period: 10 - val_period: 10 - num_preds: 10 - clip_length: 750 - window_size: 30 - s_ratio: 2 - train_eeg_head_only: False - pretrained_decoder_checkpoint: "" - pretrained_prior_checkpoint: "" - pretrained_load_prior: True - eval_eeg: True - bidirectional: False - -renderer: # face animation based on predicted 3DMM coefficients - name: PIRenderer - _target_: utils.render.Render - do_render: False - -defaults: - - /model/motion_diffusion@model - - /model/optim/adamw@optim - - /model/scheduler/cosine_annealing@scheduler - - /model/losses/motion_diffusion@criterion # loss func diff --git a/personalised/code/configs/trainer/motion_diffusion_velocity.yaml b/personalised/code/configs/trainer/motion_diffusion_velocity.yaml deleted file mode 100644 index 7d2a62f40ec2b46642c093b78ce02b99b20a1f67..0000000000000000000000000000000000000000 --- a/personalised/code/configs/trainer/motion_diffusion_velocity.yaml +++ /dev/null @@ -1,48 +0,0 @@ -model_name: motion_diffusion_velocity -_target_: trainer.motion_diffusion.Trainer - -trainer_mode: generic # generic facial reaction generator -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} # true | false -current_runid: ${run_id} -resume_runid: ${resume_id} -batch_size: 4 -post_config_name: configs/model/emotion_autoencoder.yaml -post_clip_length: 1000 # ground-truths alignment -data_clamp: True - -generic: - seed: 1234 - start_epoch: 0 - epochs: 500 - lr: 0.0001 - model: LatentMatcher - clip_grad: False - num_workers: 16 - log_dir: log - tb_dir: tb_logs - out_dir: results - save_period: 10 - val_period: 10 - num_preds: 10 - clip_length: 750 - window_size: 30 - s_ratio: 2 - train_eeg_head_only: False - pretrained_decoder_checkpoint: "" - pretrained_prior_checkpoint: "" - pretrained_load_prior: True - eval_eeg: True - bidirectional: False - -renderer: # face animation based on predicted 3DMM coefficients - name: PIRenderer - _target_: utils.render.Render - do_render: False - -defaults: - - /model/motion_diffusion_velocity@model - - /model/optim/adamw@optim - - /model/scheduler/cosine_annealing@scheduler - - /model/losses/motion_diffusion@criterion # loss func diff --git a/personalised/code/configs/trainer/motion_transvae.yaml b/personalised/code/configs/trainer/motion_transvae.yaml deleted file mode 100644 index 3069fac3d0e465f1733827a8953abdb270a425bb..0000000000000000000000000000000000000000 --- a/personalised/code/configs/trainer/motion_transvae.yaml +++ /dev/null @@ -1,42 +0,0 @@ -model_name: motion_transvae -_target_: trainer.motion_transvae.Trainer - -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} # true | false -current_runid: ${run_id} -resume_runid: ${resume_id} - -epochs: 100 -gpu_ids: 0 -lr: 0.00001 -j: 12 # num_workers -batch_size: 4 -max_seq_len: 750 -window_size: 8 -div_p: 10 -post_config_name: configs/model/emotion_autoencoder.yaml -post_clip_length: 1000 # ground-truths alignment -data_clamp: True -data_transform: zero_center # standard | zero_center -train_eeg_head_only: False -pretrained_model_checkpoint: "" -eval_eeg: True -num_preds: 10 -bidirectional: False -save_results: True -eval_facial_metrics: True -eval_eeg_metrics: True -metric_threads: 1 -eval_clip_batch_size: 1 - -renderer: # face animation based on predicted 3DMM coefficients - name: PIRenderer - _target_: utils.render.Render - do_render: True - transform_reverse: ${trainer.data_transform} - -defaults: - - /model/motion_transvae@model - - /model/optim/adamw@optim - - /model/losses/motion_transvae@criterion # loss func diff --git a/personalised/code/configs/trainer/perfrdiff_rewrite_weight.yaml b/personalised/code/configs/trainer/perfrdiff_rewrite_weight.yaml deleted file mode 100644 index 317d04df69f2b14725b5c5d91ec528d88cd34509..0000000000000000000000000000000000000000 --- a/personalised/code/configs/trainer/perfrdiff_rewrite_weight.yaml +++ /dev/null @@ -1,89 +0,0 @@ -model_name: perfrdiff_rewrite_weight -_target_: trainer.perfrdiff_rewrite_weight.Trainer - -trainer_mode: generic -folder: save/${trainer.model_name}/${data.data_name}/${task} -ckpt_dir: ${get_last_checkpoint:${trainer.folder}} -resumed_training: ${resume} -current_runid: ${run_id} -resume_runid: ${resume_id} -batch_size: 4 -post_config_name: configs/model/emotion_autoencoder.yaml -post_clip_length: 1000 -data_clamp: True -num_eval_preds: 10 -eval_clip_batch_size: 16 -parallel_eval_preds: True - -pretrained: - diffusion_prior: pretrained_models/diffusion_model/DiffusionPriorNetwork/checkpoint.pth - diffusion_decoder: pretrained_models/diffusion_model/TransformerDenoiser/checkpoint.pth - modifier_checkpoint: "" - eeg_head_checkpoint: pretrained_models/rewrite_weight/EEGPredictionHead/checkpoint.pth - -person_specific: - checkpoint_path: pretrained_models/person_specific/checkpoint.pth - args: - in_features: 58 - embed_dim: 512 - num_heads: 4 - num_layers: 4 - mlp_dim: 1024 - seq_len: 750 - proj_dim: 512 - proj_head: mlp - drop_prob: 0.1 - max_len: 1000 - pos_encoding: absolute - embed_layer: linear - -main_model: - args: - input_dim: 512 - latent_dim: 1024 - embed_dim: 512 - personal_condition_mode: 3dmm_personality - personality_input_dim: 5 - personality_hidden_dim: 128 - personality_dropout: 0.1 - personality_fusion_hidden_dim: 512 - regularization: False - regular_w: 0.0 - num_shared_layers: 2 - modified_layers: - - diffusion_decoder.model.decoder.layers.4.multihead_attn - - diffusion_decoder.model.decoder.layers.6.multihead_attn - - diffusion_decoder.model.to_emotion_feat - predict: shift - modify: all - optimizer_hypernet: - type: sgd - args: - lr: 0.001 - weight_decay: 1e-4 - momentum: 0.9 - -generic: - seed: 1234 - start_epoch: 0 - epochs: 100 - model: LatentMatcher - clip_grad: False - num_workers: 16 - log_dir: log - tb_dir: tb_logs - out_dir: results - save_period: 10 - val_period: 10 - num_preds: 1 - clip_length: 750 - window_size: 30 - s_ratio: 2 - train_eeg: False - train_eeg_head_only: False - eval_eeg: True - bidirectional: False - -defaults: - - /model/motion_diffusion@model - - /model/losses/perfrdiff_rewrite_weight@criterion diff --git a/personalised/code/dataset/__init__.py b/personalised/code/dataset/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/personalised/code/dataset/data_preprocess/audio_feature_extraction.py b/personalised/code/dataset/data_preprocess/audio_feature_extraction.py deleted file mode 100644 index bb3f7c94f0a36399dc6f81ce3019d35f84248862..0000000000000000000000000000000000000000 --- a/personalised/code/dataset/data_preprocess/audio_feature_extraction.py +++ /dev/null @@ -1,76 +0,0 @@ -import argparse -import os -import cv2 -import numpy as np -import torch -from tqdm import tqdm -from dataset.modules.audio_processor import AudioProcessor - - -def main(args): - video_paths = [] - input_paths = [] - output_paths = [] - - for root, _, files in os.walk(args.root_dir): - for file in files: - if not file.endswith(".wav"): - continue - - input_dir = os.path.join(root, file) - input_paths.append(input_dir) - - output_dir = root.replace("audio", "audio-features") - os.makedirs(output_dir, exist_ok=True) - output_path = os.path.join(output_dir, file.replace(".wav", ".npy")) - output_paths.append(output_path) - - video_dir = root.replace("audio", "video-raw") - video_path = os.path.join(video_dir, file.replace(".wav", ".mp4")) - video_paths.append(video_path) - print(f"Read {len(input_paths)} audio files, {len(output_paths)} output files, {len(video_paths)} video files") - - sample_rate = args.sample_rate - fps = args.fps - wav2vec_model_path = args.wav2vec_model_path - wav2vec_only_last_features = args.wav2vec_only_last_features == "last" - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - with AudioProcessor( - sample_rate, - fps, - wav2vec_model_path, - wav2vec_only_last_features, - device=device, - ) as audio_processor: - - for input_path, output_path, video_path in tqdm(zip(input_paths, output_paths, video_paths)): - cap = cv2.VideoCapture(video_path) - seq_len = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - - audio_emb, audio_length = audio_processor.preprocess(wav_file=input_path, seq_len=seq_len) - - try: - np.save(output_path, audio_emb) - print(f"Successfully saved audio embedding {input_path} -> {output_path}") - except Exception as e: - print(f"Error saving {input_path} -> {output_path}: {e}") - continue - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Audio Feature Extraction") - parser.add_argument('--root_dir', type=str, - default='./data', - help="root directory of react2026 dataset") - parser.add_argument('--sample_rate', type=int, default=16000, - help="original sampling rate of audio data") - parser.add_argument('--fps', type=int, default=30, help="original fps of video data") - parser.add_argument('--wav2vec_model_path', type=str, - default='./pretrained_models/wav2vec/wav2vec2-base-960h', - help="wav2vec model path") - parser.add_argument('--wav2vec_only_last_features', type=str, default='last', - help="wav2vec only last features") - - args = parser.parse_args() - - main(args) \ No newline at end of file diff --git a/personalised/code/dataset/data_preprocess/audio_feature_extraction_test.py b/personalised/code/dataset/data_preprocess/audio_feature_extraction_test.py deleted file mode 100644 index 7827c031ec70e08de3e193f261f370194efdc540..0000000000000000000000000000000000000000 --- a/personalised/code/dataset/data_preprocess/audio_feature_extraction_test.py +++ /dev/null @@ -1,120 +0,0 @@ -""" -Audio feature extraction for test split. -Same logic as audio_feature_extraction.py, but: - - uses video-face-crop instead of video-raw (test set has no video-raw) - - bypasses audio_separator (not installed); uses wav2vec directly -""" -import argparse -import math -import os -import traceback - -import cv2 -import librosa -import numpy as np -import torch -from tqdm import tqdm -from transformers import Wav2Vec2FeatureExtractor - -from framework.feature_extractor.wav2vec import Wav2VecModel - - -def main(args): - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - print(f"Using device: {device}") - - audio_encoder = Wav2VecModel.from_pretrained( - args.wav2vec_model_path, local_files_only=True - ).to(device) - audio_encoder.eval() - audio_encoder.feature_extractor._freeze_parameters() - - wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained( - args.wav2vec_model_path, local_files_only=True - ) - print("Wav2Vec model loaded.") - - video_paths = [] - input_paths = [] - output_paths = [] - - for root, _, files in os.walk(args.root_dir): - for file in sorted(files): - if not file.endswith(".wav"): - continue - - input_paths.append(os.path.join(root, file)) - - # replace only the first occurrence to avoid corrupt paths - output_dir = root.replace("/audio/", "/audio-features/", 1) - os.makedirs(output_dir, exist_ok=True) - output_paths.append(os.path.join(output_dir, file.replace(".wav", ".npy"))) - - video_dir = root.replace("/audio/", "/video-face-crop/", 1) - video_paths.append(os.path.join(video_dir, file.replace(".wav", ".mp4"))) - - print(f"Found {len(input_paths)} audio files") - - saved = 0 - skipped = 0 - for input_path, output_path, video_path in tqdm( - zip(input_paths, output_paths, video_paths), total=len(input_paths) - ): - if os.path.exists(output_path): - skipped += 1 - continue - - try: - # get frame count from video-face-crop - cap = cv2.VideoCapture(video_path) - seq_len = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - cap.release() - - # fallback: compute from audio duration if OpenCV returns 0 - if seq_len == 0: - print(f"Warning: OpenCV returned 0 frames for {video_path}, computing from audio") - y_tmp, _ = librosa.load(input_path, sr=args.sample_rate, duration=1.0) - duration = librosa.get_duration(path=input_path) - seq_len = math.ceil(duration * args.fps) - - if seq_len == 0: - print(f"Error: could not determine seq_len for {input_path}, skipping") - continue - - # load audio - speech_array, sampling_rate = librosa.load(input_path, sr=args.sample_rate) - audio_feature = np.squeeze( - wav2vec_feature_extractor( - speech_array, sampling_rate=sampling_rate - ).input_values - ) - - audio_feature_t = torch.from_numpy(audio_feature).float().unsqueeze(0).to(device) - - with torch.no_grad(): - embeddings = audio_encoder( - audio_feature_t, seq_len=seq_len, output_hidden_states=True - ) - - audio_emb = embeddings.last_hidden_state.squeeze().cpu().detach() # tensor - - np.save(output_path, audio_emb) - saved += 1 - - except Exception: - print(f"Error processing {input_path}:") - traceback.print_exc() - - print(f"Done. Saved: {saved}, Skipped (already exist): {skipped}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Audio Feature Extraction for Test Split") - parser.add_argument("--root_dir", type=str, - default="/mnt/HDD1/MARS/test/audio/speaker") - parser.add_argument("--sample_rate", type=int, default=16000) - parser.add_argument("--fps", type=int, default=30) - parser.add_argument("--wav2vec_model_path", type=str, - default="./pretrained_models/wav2vec/wav2vec2-base-960h") - args = parser.parse_args() - main(args) diff --git a/personalised/code/dataset/modules/audio_embedder.py b/personalised/code/dataset/modules/audio_embedder.py deleted file mode 100644 index 792207f3bf7bce9950528383b339081af15f49ca..0000000000000000000000000000000000000000 --- a/personalised/code/dataset/modules/audio_embedder.py +++ /dev/null @@ -1,16 +0,0 @@ -import torch -import torch.nn as nn - - -class AudioEmbedder(nn.Module): - def __init__(self, skip_norm=False, *args, **kwargs): - super().__init__() - self.skip_norm = skip_norm - - def _encode(self, x): - if not self.skip_norm: - x_min = torch.min(x) - x_max = torch.max(x) - return (x - x_min) / (x_max - x_min) - else: - return x diff --git a/personalised/code/dataset/modules/audio_processor.py b/personalised/code/dataset/modules/audio_processor.py deleted file mode 100644 index faa1117ea062b6ef9510ead08e3079905fc56006..0000000000000000000000000000000000000000 --- a/personalised/code/dataset/modules/audio_processor.py +++ /dev/null @@ -1,189 +0,0 @@ -# pylint: disable=C0301 -''' -This module contains the AudioProcessor class and related functions for processing audio data. -It utilizes various libraries and models to perform tasks such as preprocessing, feature extraction, -and audio separation. The class is initialized with configuration parameters and can process -audio files using the provided models. -''' -import math -import os -import librosa -import numpy as np -import torch -from audio_separator.separator import Separator -from einops import rearrange -from transformers import Wav2Vec2FeatureExtractor -from dataset.tools.util import resample_audio -from framework.feature_extractor.wav2vec import Wav2VecModel -import soundfile as sf -import torchaudio - -torchaudio.set_audio_backend("sox_io") - - -class AudioProcessor: - """ - AudioProcessor is a class that handles the processing of audio files. - It takes care of preprocessing the audio files, extracting features - using wav2vec models, and separating audio signals if needed. - - :param sample_rate: Sampling rate of the audio file - :param fps: Frames per second for the extracted features - :param wav2vec_model_path: Path to the wav2vec model - :param only_last_features: Whether to only use the last features - :param audio_separator_model_path: Path to the audio separator model - :param audio_separator_model_name: Name of the audio separator model - :param cache_dir: Directory to cache the intermediate results - :param device: Device to run the processing on - """ - - def __init__( - self, - sample_rate, - fps, - wav2vec_model_path, - only_last_features: bool = True, - audio_separator_model_path: str = None, - audio_separator_model_name: str = None, - cache_dir: str = '', - device="cuda:0", - ) -> None: - self.sample_rate = sample_rate - self.fps = fps - self.device = device - - self.audio_encoder = Wav2VecModel.from_pretrained(wav2vec_model_path, local_files_only=True).to(device=device) - self.audio_encoder.feature_extractor._freeze_parameters() - self.only_last_features = only_last_features - - if audio_separator_model_name is not None: - try: - os.makedirs(cache_dir, exist_ok=True) - except OSError as _: - print("Fail to create the output cache dir.") - self.audio_separator = Separator( - output_dir=cache_dir, - output_single_stem="vocals", - model_file_dir=audio_separator_model_path, - ) - self.audio_separator.load_model(audio_separator_model_name) - assert self.audio_separator.model_instance is not None, "Fail to load audio separate model." - else: - self.audio_separator = None - print("Use audio directly without vocals seperator.") - - self.wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec_model_path, - local_files_only=True) - - def preprocess(self, wav_file: str, clip_length: int = -1, seq_len: int = None): - """ - Preprocess a WAV audio file by separating the vocals from the background and resampling it to a 16 kHz sample rate. - The separated vocal track is then converted into wav2vec2 for further processing or analysis. - - Args: - wav_file (str): The path to the WAV file to be processed. This file should be accessible and in WAV format. - - Raises: - RuntimeError: Raises an exception if the WAV file cannot be processed. This could be due to issues - such as file not found, unsupported file format, or errors during the audio processing steps. - - Returns: - torch.tensor: Returns an audio embedding as a torch.tensor - """ - if self.audio_separator is not None: - # 1. separate vocals - outputs = self.audio_separator.separate(wav_file) - if len(outputs) <= 0: - raise RuntimeError("Audio separate failed.") - - vocal_audio_file = outputs[0] - vocal_audio_name, _ = os.path.splitext(vocal_audio_file) - vocal_audio_file = os.path.join(self.audio_separator.output_dir, vocal_audio_file) - vocal_audio_file = resample_audio(vocal_audio_file, - os.path.join(self.audio_separator.output_dir, - f"{vocal_audio_name}-16k.wav"), - self.sample_rate) - else: - vocal_audio_file = wav_file - - # 2. extract wav2vec features - speech_array, sampling_rate = librosa.load(vocal_audio_file, sr=self.sample_rate) - audio_feature = np.squeeze( - self.wav2vec_feature_extractor(speech_array, sampling_rate=sampling_rate).input_values) - # compute sequence length based on video fps - if seq_len is None: - seq_len = math.ceil(len(audio_feature) / self.sample_rate * self.fps) - else: - if math.ceil(len(audio_feature) / self.sample_rate * self.fps) != seq_len: - print(f"Warning: seq_len {seq_len} is not equal to " - f"computed version {math.ceil(len(audio_feature) / self.sample_rate * self.fps)}") - audio_length = seq_len - - audio_feature = torch.from_numpy(audio_feature).float().to(device=self.device) - - if clip_length > 0 and seq_len % clip_length != 0: - audio_feature = torch.nn.functional.pad(audio_feature, (0, (clip_length - seq_len % clip_length) * ( - self.sample_rate // self.fps)), 'constant', 0.0) - seq_len += clip_length - seq_len % clip_length - audio_feature = audio_feature.unsqueeze(0) - - with torch.no_grad(): - embeddings = self.audio_encoder(audio_feature, seq_len=seq_len, output_hidden_states=True) - assert len(embeddings) > 0, "Fail to extract audio embedding" - if self.only_last_features: - audio_emb = embeddings.last_hidden_state.squeeze() - else: - audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0) - audio_emb = rearrange(audio_emb, "b s d -> s b d") - - audio_emb = audio_emb.cpu().detach() - - return audio_emb, audio_length - - def get_embedding(self, wav_file: str): - """preprocess wav audio file convert to embeddings - - Args: - wav_file (str): The path to the WAV file to be processed. This file should be accessible and in WAV format. - - Returns: - torch.tensor: Returns an audio embedding as a torch.tensor - """ - speech_array, sampling_rate = librosa.load( - wav_file, sr=self.sample_rate) - assert sampling_rate == 16000, "The audio sample rate must be 16000" - audio_feature = np.squeeze(self.wav2vec_feature_extractor( - speech_array, sampling_rate=sampling_rate).input_values) - seq_len = math.ceil(len(audio_feature) / self.sample_rate * self.fps) - - audio_feature = torch.from_numpy( - audio_feature).float().to(device=self.device) - audio_feature = audio_feature.unsqueeze(0) - - with torch.no_grad(): - embeddings = self.audio_encoder( - audio_feature, seq_len=seq_len, output_hidden_states=True) - assert len(embeddings) > 0, "Fail to extract audio embedding" - - if self.only_last_features: - audio_emb = embeddings.last_hidden_state.squeeze() - else: - audio_emb = torch.stack( - embeddings.hidden_states[1:], dim=1).squeeze(0) - audio_emb = rearrange(audio_emb, "b s d -> s b d") - - audio_emb = audio_emb.cpu().detach() - - return audio_emb - - def close(self): - """ - TODO: to be implemented - """ - return self - - def __enter__(self): - return self - - def __exit__(self, _exc_type, _exc_val, _exc_tb): - self.close() diff --git a/personalised/code/dataset/perfrdiff_rewrite_weight.py b/personalised/code/dataset/perfrdiff_rewrite_weight.py deleted file mode 100644 index 7a9062b2531fd86d9ef120670397512d3c593739..0000000000000000000000000000000000000000 --- a/personalised/code/dataset/perfrdiff_rewrite_weight.py +++ /dev/null @@ -1,534 +0,0 @@ -import csv -import io -import os -import random -import tarfile -from pathlib import Path - -import hydra -import numpy as np -import torch -from hydra.utils import instantiate -from omegaconf import DictConfig -from torch.utils.data import DataLoader, Dataset -from torchvision import transforms - -from dataset.tools.util import extract_audio_features, Transform -from decord import VideoReader, cpu -from PIL import Image - - -DEFAULT_EEG_TARGET_COLS = [ - "TP9", "AF7", "AF8", "TP10", - "Delta_TP9", "Theta_TP9", "Alpha_TP9", "Beta_TP9", "Gamma_TP9", - "Delta_TP10", "Theta_TP10", "Alpha_TP10", "Beta_TP10", "Gamma_TP10", -] -EEG_RAW_CHANNELS = {"TP9", "AF7", "AF8", "TP10"} - - -def _empty_tensor(): - return torch.zeros(size=(0,)) - - -def collate_fit(batch): - columns = list(zip(*batch)) - return tuple(torch.stack(items, dim=0) if items[0].numel() > 0 else torch.zeros(size=(len(items), 0)) - for items in columns) - - -def collate_test(batch): - columns = list(zip(*batch)) - return tuple(list(column) for column in columns) - - -class PerFRDiffRewriteWeightDataset(Dataset): - def __init__( - self, - root_dir, - split="train", - clip_length=750, - target_size=224, - crop_size=224, - fps=30, - audio_feature_type="wav2vec", - load_video_s=False, - load_video_l=False, - load_audio=True, - load_emotion_s=True, - load_emotion_l=True, - load_3dmm_s=True, - load_3dmm_l=True, - load_ref=False, - load_personality_l=False, - personal_condition_mode="3dmm_personality", - personality_dir_name="personality", - load_eeg_l=False, - eeg_dir_name="eeg_processed", - eeg_target_cols=None, - eeg_channel_scale=1000.0, - eeg_use_tar_fallback=True, - normalize_3dmm="standard", - num_test_gts=10, - bidirectional=False, - **kwargs, - ): - self.root_dir = root_dir - self.split = split - self.clip_length = clip_length - self.fps = fps - self.audio_feature_type = audio_feature_type - self.load_video_s = load_video_s - self.load_video_l = load_video_l - self.load_audio = load_audio - self.load_emotion_s = load_emotion_s - self.load_emotion_l = load_emotion_l - self.load_3dmm_s = load_3dmm_s - self.load_3dmm_l = load_3dmm_l - self.load_ref = load_ref - self.personal_condition_mode = personal_condition_mode - if self.personal_condition_mode not in {"3dmm_personality", "personality_only", "3dmm_only"}: - raise ValueError(f"Unknown personal_condition_mode: {self.personal_condition_mode}") - self.load_personality_l = load_personality_l and self.personal_condition_mode != "3dmm_only" - self.load_eeg_l = load_eeg_l - self.eeg_target_cols = list(eeg_target_cols) if eeg_target_cols is not None else DEFAULT_EEG_TARGET_COLS - self.eeg_channel_scale = eeg_channel_scale - self.eeg_use_tar_fallback = eeg_use_tar_fallback - self._eeg_tar_members = None - self.num_test_gts = num_test_gts - self.bidirectional = bidirectional - - split_dir = Path(root_dir) / split - self.audio_dir = split_dir / ("audio-features" if audio_feature_type == "wav2vec" else "audio") - self.video_dir = split_dir / "video-face-crop" - # (FRRea) listener video loading for FID rendering: only active when load_video_l. - # frrea_video_stride>0 decodes only every Nth frame (frame 0 = render reference). - self._frrea_video_stride = int(kwargs.get("frrea_video_stride", 0)) - self._video_transform = Transform(target_size, crop_size) - self.emotion_dir = split_dir / "facial-attributes" - self.coeff_dir = split_dir / "coefficients" - self.personality_dir = split_dir / personality_dir_name - self.eeg_dir = split_dir / eeg_dir_name - self.eeg_tar_path = split_dir / f"{eeg_dir_name}.tar.gz" - self.personality_cols = [ - "Extraversion", - "Agreeableness", - "Conscientiousness", - "Neuroticism", - "Openness", - ] - self.listener_personality_by_session = {} - self.listener_personality_by_stem = {} - self.personality_by_role_session = {} - self.personality_by_role_stem = {} - if self.load_personality_l: - self._load_personality_index("listener", required=True) - if self.bidirectional: - self._load_personality_index("speaker", required=False) - - mean_face_path = os.path.join(hydra.utils.get_original_cwd(), "external/FaceVerse/mean_face.npy") - std_face_path = os.path.join(hydra.utils.get_original_cwd(), "external/FaceVerse/std_face.npy") - self.mean_face = torch.FloatTensor(np.load(mean_face_path).astype(np.float32)).view(1, 1, -1) - self.std_face = torch.FloatTensor(np.load(std_face_path).astype(np.float32)).view(1, 1, -1) - if normalize_3dmm == "standard": - self.transform_3dmm = transforms.Lambda(lambda e: (e - self.mean_face) / self.std_face) - elif normalize_3dmm == "zero_center": - self.transform_3dmm = transforms.Lambda(lambda e: e - self.mean_face) - else: - raise ValueError(f"Unknown normalize_3dmm: {normalize_3dmm}") - - self.samples = [] - self.gt_path_dict = self._build_gt_index() - self._build_samples() - - @staticmethod - def _is_real_sample_file(file_name): - path = Path(file_name) - if path.suffix.lower() != ".npy": - return False - return not path.name.startswith(".") and not path.name.startswith("._") - - def _iter_emotion_files(self): - for root, dirs, files in os.walk(self.emotion_dir): - dirs[:] = [ - name for name in dirs - if not name.startswith(".") and name not in {"__MACOSX", "PaxHeader"} - ] - root_parts = Path(root).parts - if len(root_parts) < 2: - continue - role = root_parts[-2] - session = root_parts[-1] - if role not in {"speaker", "listener"}: - continue - for file_name in files: - if not self._is_real_sample_file(file_name): - continue - yield role, session, Path(file_name).stem - - def _build_gt_index(self): - gt_path_dict = {} - for role, session, stem in self._iter_emotion_files(): - session_id = Path(role) / session - file_path = session_id / stem - gt_path_dict.setdefault(session_id, []).append(file_path) - return gt_path_dict - - def _load_personality_index(self, role, required=True): - path = self.personality_dir / f"{role}.csv" - if not path.exists(): - if required: - raise FileNotFoundError(f"Missing {role} personality file: {path}") - return - - with path.open("r", newline="", encoding="utf-8-sig") as file: - reader = csv.DictReader(file) - for row in reader: - stem = Path(row.get("video_name", "")).stem - if not stem: - continue - values = [ - (float(row[column]) - 1.0) / 4.0 - for column in self.personality_cols - ] - personality = torch.tensor(values, dtype=torch.float32) - session = row.get("session") - if session: - self.personality_by_role_session[(role, session, stem)] = personality - if role == "listener": - self.listener_personality_by_session[(session, stem)] = personality - self.personality_by_role_stem[(role, stem)] = personality - if role == "listener": - self.listener_personality_by_stem[stem] = personality - - def _has_target_personality(self, listener_path): - if not self.load_personality_l: - return True - - role = listener_path.parts[0] - session = listener_path.parts[1] if len(listener_path.parts) > 1 else None - stem = listener_path.stem - if session is not None and (role, session, stem) in self.personality_by_role_session: - return True - return (role, stem) in self.personality_by_role_stem - - def _has_required(self, speaker_path, listener_path): - required = [ - self.emotion_dir / speaker_path.with_suffix(".npy"), - self.emotion_dir / listener_path.with_suffix(".npy"), - self.coeff_dir / speaker_path.with_suffix(".npy"), - self.coeff_dir / listener_path.with_suffix(".npy"), - ] - if self.audio_feature_type == "wav2vec": - required.append(self.audio_dir / speaker_path.with_suffix(".npy")) - else: - required.append(self.audio_dir / speaker_path.with_suffix(".wav")) - return all(path.exists() for path in required) - - def _build_samples(self): - for role, session, stem in self._iter_emotion_files(): - if not self.bidirectional and role != "speaker": - continue - speaker_path = Path(role) / session / stem - listener_role = "listener" if role == "speaker" else "speaker" - listener_path = Path(listener_role) / session / stem - if not self._has_required(speaker_path, listener_path): - continue - if not self._has_target_personality(listener_path): - continue - self.samples.append( - { - "speaker_path": speaker_path, - "listener_path": listener_path, - "gt_paths": self.gt_path_dict.get(Path(listener_role) / session, [listener_path]), - } - ) - - def __len__(self): - return len(self.samples) - - def _pad_clip(self, clip, target_len): - clip = clip[:target_len] - if clip.shape[0] >= target_len: - return clip - pad_shape = (target_len - clip.shape[0], *clip.shape[1:]) - return torch.cat((clip, clip.new_zeros(pad_shape)), dim=0) - - @staticmethod - def _load_numpy_array(path, name): - try: - return np.load(path, allow_pickle=False) - except ValueError as exc: - if "pickled" in str(exc): - raise ValueError( - f"{name} file is not a numeric numpy sample: {path}. " - "It is probably an AppleDouble/archive metadata file such as '._*.npy'. " - "Remove those metadata files from the dataset or keep them hidden so the loader can skip them." - ) from exc - raise - - def _load_emotion(self, rel_path): - path = self.emotion_dir / rel_path.with_suffix(".npy") - return torch.from_numpy(self._load_numpy_array(path, "emotion")).float() - - def _load_audio(self, rel_path, total_length): - if self.audio_feature_type == "wav2vec": - path = self.audio_dir / rel_path.with_suffix(".npy") - audio = self._load_numpy_array(path, "audio") - elif self.audio_feature_type == "mfcc": - audio = extract_audio_features(os.fspath(self.audio_dir / rel_path.with_suffix(".wav")), - self.fps, - total_length) - else: - raise ValueError(f"Unknown audio_feature_type: {self.audio_feature_type}") - return torch.from_numpy(audio).float() - - def _load_3dmm(self, rel_path): - path = self.coeff_dir / rel_path.with_suffix(".npy") - coeff = torch.FloatTensor(self._load_numpy_array(path, "3DMM")).squeeze() - return self.transform_3dmm(coeff)[0].float() - - def _load_personality(self, listener_path): - if not self.load_personality_l: - return _empty_tensor() - - role = listener_path.parts[0] - session = listener_path.parts[1] if len(listener_path.parts) > 1 else None - stem = listener_path.stem - if session is not None and (role, session, stem) in self.personality_by_role_session: - return self.personality_by_role_session[(role, session, stem)].clone() - if (role, stem) in self.personality_by_role_stem: - return self.personality_by_role_stem[(role, stem)].clone() - raise FileNotFoundError(f"Missing target personality for {listener_path}") - - @staticmethod - def _is_archive_metadata(path): - parts = Path(path).parts - name = Path(path).name - return name.startswith("._") or name == ".DS_Store" or "PaxHeader" in parts - - def _read_eeg_text(self, rel_path): - eeg_path = self.eeg_dir / rel_path.with_suffix(".csv") - if eeg_path.exists() and not self._is_archive_metadata(eeg_path): - return eeg_path.read_text(encoding="utf-8-sig") - - if not self.eeg_use_tar_fallback or not self.eeg_tar_path.exists(): - return None - - if self._eeg_tar_members is None: - with tarfile.open(self.eeg_tar_path, "r:gz") as tar: - self._eeg_tar_members = { - name for name in tar.getnames() - if not self._is_archive_metadata(name) - } - - member_candidates = [ - os.fspath(Path(self.eeg_dir.name) / rel_path.with_suffix(".csv")).replace("\\", "/"), - os.fspath(rel_path.with_suffix(".csv")).replace("\\", "/"), - ] - member = next((name for name in member_candidates if name in self._eeg_tar_members), None) - if member is None: - return None - - with tarfile.open(self.eeg_tar_path, "r:gz") as tar: - extracted = tar.extractfile(member) - if extracted is None: - return None - return extracted.read().decode("utf-8-sig") - - def _load_eeg(self, listener_path, total_length): - eeg_dim = len(self.eeg_target_cols) - empty_target = torch.zeros(size=(total_length, eeg_dim), dtype=torch.float32) - empty_mask = torch.zeros(size=(total_length, eeg_dim), dtype=torch.float32) - - text = self._read_eeg_text(listener_path) - if text is None: - return empty_target, empty_mask - - rows = list(csv.DictReader(io.StringIO(text))) - if len(rows) == 0: - return empty_target, empty_mask - - values = np.zeros((len(rows), eeg_dim), dtype=np.float32) - mask = np.zeros((len(rows), eeg_dim), dtype=np.float32) - for row_idx, row in enumerate(rows): - for col_idx, col in enumerate(self.eeg_target_cols): - raw_value = row.get(col, "") - if raw_value == "": - continue - try: - value = float(raw_value) - except ValueError: - continue - if not np.isfinite(value): - continue - if col in EEG_RAW_CHANNELS: - value = value / self.eeg_channel_scale - values[row_idx, col_idx] = value - mask[row_idx, col_idx] = 1.0 - - frame_to_eeg = np.floor(np.arange(total_length) / self.fps).astype(np.int64) - frame_to_eeg = np.clip(frame_to_eeg, 0, len(rows) - 1) - return ( - torch.from_numpy(values[frame_to_eeg]), - torch.from_numpy(mask[frame_to_eeg]), - ) - - def _choose_personal_path(self, listener_path): - session_id = Path(*listener_path.parts[:2]) - candidates = [path for path in self.gt_path_dict.get(session_id, []) if path != listener_path] - if not candidates: - return listener_path - return random.choice(candidates) - - def _sample_test_gts(self, listener_path, gt_paths): - candidates = [path for path in gt_paths if path != listener_path] - paths = [listener_path] - if len(candidates) >= self.num_test_gts - 1: - paths.extend(random.sample(candidates, self.num_test_gts - 1)) - elif candidates: - paths.extend(random.choices(candidates, k=self.num_test_gts - 1)) - else: - paths.extend([listener_path] * (self.num_test_gts - 1)) - return paths - - def _load_listener_video(self, rel_path): - """(FRRea) Load the GT listener face video for FID rendering. Returns a - (M, 3, H, W) normalized clip; frame 0 doubles as the render reference.""" - video_path = os.fspath(self.video_dir / rel_path.with_suffix(".mp4")) - vr = VideoReader(video_path, ctx=cpu(0)) - if self._frrea_video_stride and self._frrea_video_stride > 0: - indices = list(range(0, len(vr), self._frrea_video_stride)) - frames = vr.get_batch(indices).asnumpy() - clip = [self._video_transform(Image.fromarray(frames[j])) for j in range(frames.shape[0])] - else: - clip = [self._video_transform(Image.fromarray(vr[f].asnumpy())) for f in range(len(vr))] - del vr - return torch.stack(clip, dim=0) - - def _load_personal_clip(self, listener_path, deterministic=False): - personal_3dmm = self._load_3dmm(self._choose_personal_path(listener_path)) - if deterministic or personal_3dmm.shape[0] <= self.clip_length: - cp = 0 - else: - cp = random.randint(0, personal_3dmm.shape[0] - self.clip_length) - return self._pad_clip(personal_3dmm[cp:cp + self.clip_length], self.clip_length) - - def __getitem__(self, index): - sample = self.samples[index] - speaker_path = sample["speaker_path"] - listener_path = sample["listener_path"] - - speaker_emotion = self._load_emotion(speaker_path) - listener_emotion = self._load_emotion(listener_path) - speaker_3dmm = self._load_3dmm(speaker_path) - listener_3dmm = self._load_3dmm(listener_path) - listener_personality = self._load_personality(listener_path) - - total_length = min( - speaker_emotion.shape[0], - listener_emotion.shape[0], - speaker_3dmm.shape[0], - listener_3dmm.shape[0], - ) - speaker_audio = self._load_audio(speaker_path, total_length) - total_length = min(total_length, speaker_audio.shape[0]) - listener_eeg = listener_eeg_mask = _empty_tensor() - if self.load_eeg_l: - listener_eeg, listener_eeg_mask = self._load_eeg(listener_path, total_length) - - if self.split == "test": - gt_paths = self._sample_test_gts(listener_path, sample["gt_paths"]) - listener_emotion_gts = [self._load_emotion(path) for path in gt_paths] - listener_3dmm_gts = [self._load_3dmm(path) for path in gt_paths] - listener_clip_lengths = torch.tensor([emotion.shape[0] for emotion in listener_emotion_gts]) - personal_3dmm = _empty_tensor() if self.personal_condition_mode == "personality_only" \ - else self._load_personal_clip(listener_path, deterministic=True) - # (FRRea) GT listener video at slot 4 (else empty); gt_paths[0] is the target listener. - listener_video = self._load_listener_video(gt_paths[0]) if self.load_video_l else _empty_tensor() - return ( - speaker_audio[:total_length], - _empty_tensor(), - speaker_emotion[:total_length], - speaker_3dmm[:total_length], - listener_video, - listener_emotion_gts, - listener_3dmm_gts, - personal_3dmm, - listener_personality, - listener_eeg, - listener_eeg_mask, - torch.tensor(total_length), - listener_clip_lengths, - ) - - cp = random.randint(0, total_length - self.clip_length) if total_length > self.clip_length else 0 - personal_3dmm = _empty_tensor() if self.personal_condition_mode == "personality_only" \ - else self._load_personal_clip(listener_path) - - return ( - self._pad_clip(speaker_audio[cp:cp + self.clip_length], self.clip_length), - _empty_tensor(), - self._pad_clip(speaker_emotion[cp:cp + self.clip_length], self.clip_length), - self._pad_clip(speaker_3dmm[cp:cp + self.clip_length], self.clip_length), - _empty_tensor(), - self._pad_clip(listener_emotion[cp:cp + self.clip_length], self.clip_length), - self._pad_clip(listener_3dmm[cp:cp + self.clip_length], self.clip_length), - personal_3dmm, - listener_personality, - self._pad_clip(listener_eeg[cp:cp + self.clip_length], self.clip_length) - if self.load_eeg_l else listener_eeg, - self._pad_clip(listener_eeg_mask[cp:cp + self.clip_length], self.clip_length) - if self.load_eeg_l else listener_eeg_mask, - _empty_tensor(), - ) - - -class PerFRDiffRewriteWeightDataModule: - def __init__( - self, - train_dataset: DictConfig = None, - validation_dataset: DictConfig = None, - test_dataset: DictConfig = None, - **kwargs, - ): - self.seed = kwargs.pop("seed", 1234) - self.train_set_cfg = train_dataset - self.val_set_cfg = validation_dataset - self.test_set_cfg = test_dataset - - def get_dataloader(self, stage): - def worker_init_fn(worker_id): - seed = self.seed + worker_id - random.seed(seed) - np.random.seed(seed) - - if stage == "fit": - train_dataset = instantiate(self.train_set_cfg) - train_loader = DataLoader( - dataset=train_dataset, - batch_size=self.train_set_cfg.batch_size, - shuffle=self.train_set_cfg.shuffle, - num_workers=self.train_set_cfg.num_workers, - collate_fn=collate_fit, - worker_init_fn=worker_init_fn, - ) - val_dataset = instantiate(self.val_set_cfg) - val_loader = DataLoader( - dataset=val_dataset, - batch_size=self.val_set_cfg.batch_size, - shuffle=self.val_set_cfg.shuffle, - num_workers=self.val_set_cfg.num_workers, - collate_fn=collate_fit, - worker_init_fn=worker_init_fn, - ) - return train_loader, val_loader - - test_dataset = instantiate(self.test_set_cfg) - return DataLoader( - dataset=test_dataset, - batch_size=self.test_set_cfg.batch_size, - shuffle=self.test_set_cfg.shuffle, - num_workers=self.test_set_cfg.num_workers, - collate_fn=collate_test, - worker_init_fn=worker_init_fn, - ) diff --git a/personalised/code/dataset/react_2025.py b/personalised/code/dataset/react_2025.py deleted file mode 100644 index 713818969f3ba13bec3c688f115939266f9232f5..0000000000000000000000000000000000000000 --- a/personalised/code/dataset/react_2025.py +++ /dev/null @@ -1,621 +0,0 @@ -import os -import csv -import io -import tarfile -from pathlib import Path -import hydra -from hydra.utils import instantiate -from omegaconf import DictConfig -import torch -from torch.utils import data -from torchvision import transforms -import numpy as np -import random -from PIL import Image -from decord import VideoReader -from decord import cpu -from torch.utils.data import DataLoader -from dataset.tools.util import Transform, extract_audio_features -import torchaudio - -torchaudio.set_audio_backend("sox_io") - - -DEFAULT_EEG_TARGET_COLS = [ - "TP9", "AF7", "AF8", "TP10", - "Delta_TP9", "Theta_TP9", "Alpha_TP9", "Beta_TP9", "Gamma_TP9", - "Delta_TP10", "Theta_TP10", "Alpha_TP10", "Beta_TP10", "Gamma_TP10" -] -EEG_RAW_CHANNELS = {"TP9", "AF7", "AF8", "TP10"} - - -def custom_collate(batch): - speaker_audio_clip = [item[0] for item in batch if len(item[0]) > 0] - speaker_video_clip = [item[1] for item in batch if len(item[1]) > 0] - speaker_emotion_clip = [item[2] for item in batch if len(item[2]) > 0] - speaker_params_clip = [item[3] for item in batch if len(item[3]) > 0] - listener_video_clip = [item[4] for item in batch if len(item[4]) > 0] - listener_emotion_clip = [item[5] for item in batch if len(item[5]) > 0] - listener_params_clip = [item[6] for item in batch if len(item[6]) > 0] - speaker_clip_length = torch.stack([item[7] if isinstance(item[7], torch.Tensor) else torch.tensor(item[7]) for item in batch]) - listener_clip_length = torch.stack([item[8] if isinstance(item[8], torch.Tensor) else torch.tensor(item[8]) for item in batch]) - has_eeg = len(batch[0]) > 9 - listener_eeg_clip = [item[9] for item in batch] if has_eeg else None - listener_eeg_mask = [item[10] for item in batch] if has_eeg else None - - collated = ( - speaker_audio_clip, - speaker_video_clip, - speaker_emotion_clip, - speaker_params_clip, - listener_video_clip, - listener_emotion_clip, - listener_params_clip, - speaker_clip_length, - listener_clip_length, - ) - if has_eeg: - collated = collated + (listener_eeg_clip, listener_eeg_mask) - return collated - - -class ReactionAutoEncoderDataset(data.Dataset): - def __init__(self, - root_dir, - split='train', - clip_length: int = 1000, - **kwargs): - - self._root_dir = root_dir - self._split = split - self._clip_length = clip_length - - dataset_dir = os.path.join(root_dir, self._split) - self._emotion_dir = os.path.join(dataset_dir, 'facial-attributes') - - data_list = [] - for root, _, files in os.walk(self._emotion_dir): - for path in files: - path = Path(path) - - if path.suffix.lower() != '.npy': - continue - - file_path = os.path.join(root, path) - data_list.append(file_path) - self._data_list = data_list - - def __getitem__(self, index): - data_path = self._data_list[index] - emotion = np.load(data_path) - speaker_emotion_clip = torch.from_numpy(emotion) - total_length = len(emotion) - - global_cp = random.randint(0, total_length - self._clip_length) \ - if total_length > self._clip_length else 0 - speaker_emotion_clip = speaker_emotion_clip[global_cp: global_cp + self._clip_length] # (25-d) - - # padding_length = self._clip_length - total_length if total_length < self._clip_length else 0 - clip_length = self._clip_length if total_length >= self._clip_length else total_length - - input_length = random.randint(1, clip_length) - input_start_idx = random.randint(0, clip_length - input_length) - input_end_idx = input_start_idx + input_length - 1 - input_emotion_clip = speaker_emotion_clip[input_start_idx: input_start_idx + input_length] - input_emotion_clip = torch.cat( - (input_emotion_clip, torch.zeros(size=(self._clip_length - input_length, 25))), dim=0) - - output_length = random.randint(1, clip_length) - output_start_idx = random.randint(0, clip_length - output_length) - output_end_idx = output_start_idx + output_length - 1 - output_emotion_clip = speaker_emotion_clip[output_start_idx: output_start_idx + output_length] - output_emotion_clip = torch.cat( - (output_emotion_clip, torch.zeros(size=(self._clip_length - output_length, 25))), dim=0) - - return (input_emotion_clip, input_start_idx, input_end_idx, - output_emotion_clip, output_start_idx, output_end_idx) - - def __len__(self): - return len(self._data_list) - - -class ReactionDataset(data.Dataset): - def __init__(self, - root_dir: str = './data', # /path/to/dataset - split: str = 'train', - clip_length: int = None, - target_size: int = 224, - crop_size: int = 224, - fps: int = 30, - audio_feature_type: str = 'wav2vec', - load_video_s: bool = False, - load_video_l: bool = False, - load_audio: bool = True, - load_emotion_s: bool = True, - load_emotion_l: bool = True, - load_3dmm_s: bool = True, - load_3dmm_l: bool = True, - load_eeg_l: bool = False, - eeg_dir_name: str = 'eeg_processed', - eeg_target_cols=None, - eeg_channel_scale: float = 1000.0, - eeg_use_tar_fallback: bool = True, - bidirectional: bool = False, - normalize_3dmm: str = 'standard', # standard | zero_center - frrea_video_stride: int = 0, # >0: decode only every Nth listener frame (FRRea/FID) - **kwargs, - ): - - self._root_dir = root_dir - self._clip_length = clip_length - self._fps = fps - self._split = split - self.load_video_s = load_video_s - self.load_video_l = load_video_l - self.load_audio = load_audio - self.load_emotion_s = load_emotion_s - self.load_emotion_l = load_emotion_l - self.load_3dmm_s = load_3dmm_s - self.load_3dmm_l = load_3dmm_l - self.load_eeg_l = load_eeg_l - self.eeg_target_cols = list(eeg_target_cols) if eeg_target_cols is not None else DEFAULT_EEG_TARGET_COLS - self.eeg_channel_scale = float(eeg_channel_scale) - self.eeg_use_tar_fallback = eeg_use_tar_fallback - self.bidirectional = bidirectional - - dataset_dir = os.path.join(root_dir, self._split) - self.audio_feature_type = audio_feature_type - if self.audio_feature_type == 'wav2vec': - self._audio_dir = os.path.join(dataset_dir, 'audio-features') - elif self.audio_feature_type == 'mfcc': - self._audio_dir = os.path.join(dataset_dir, 'audio') - self._video_dir = os.path.join(dataset_dir, 'video-face-crop') - self._emotion_dir = os.path.join(dataset_dir, 'facial-attributes') - self._3dmm_dir = os.path.join(dataset_dir, 'coefficients') - self._eeg_dir = os.path.join(dataset_dir, eeg_dir_name) - self._eeg_tar_path = os.path.join(dataset_dir, f'{eeg_dir_name}.tar.gz') - self._eeg_tar_members = None - - gt_path_dict = {} - for root, _, files in os.walk(self._video_dir): - for path in files: - path = Path(path) - file, ext = path.stem, path.suffix - - if ext.lower() != '.mp4': - continue - - session_id = Path(*Path(root).parts[-2:]) # listener/session0 - - file_path = session_id / file - if session_id not in gt_path_dict: - gt_path_dict[session_id] = [file_path] - else: - gt_path_dict[session_id].append(file_path) - - speaker_path_list = [] - listener_path_list = [] - gt_path_list = [] - - for root, _, files in os.walk(self._video_dir): - for path in files: - path = Path(path) - file, ext = path.stem, path.suffix - # file, ext = os.path.splitext(path) - # e.g., 'Camera-2024-06-21-103121-103102', '.mp4' - if ext.lower() != '.mp4': - continue - - # role: listener | speaker - # session id: session* - # file: Camera-2024-06-21-103121-103102 - parts = Path(root).parts - file_path = Path(*parts[-2:]) / file - role = parts[-2] - session_id = Path(parts[-1]) - if not self.bidirectional and role != 'speaker': - continue - gt_session_id = 'speaker' / session_id if role == 'listener' else 'listener' / session_id - listener_file_path = gt_session_id / file - - speaker_path_list.append(file_path) - listener_path_list.append(listener_file_path) - listener_gt_paths = gt_path_dict[gt_session_id] - gt_path_list.append(listener_gt_paths) - - self.speaker_path_list = speaker_path_list.copy() - self.listener_path_list = listener_path_list.copy() - self.gt_path_list = gt_path_list.copy() - - # Optional test-time sharding for multi-GPU parallel evaluation. - # Enabled only via env vars (EVAL_SHARD_NUM/EVAL_SHARD_IDX); default behavior unchanged. - # All reported metrics (FRC/FRD/TLCC/smse/FRVar) are per-sample means, so evaluating - # disjoint strided shards and weighted-averaging by sample count is exactly equivalent. - shard_num = int(os.environ.get('EVAL_SHARD_NUM', '1')) - shard_idx = int(os.environ.get('EVAL_SHARD_IDX', '0')) - if shard_num > 1 and self._split == 'test': - order = sorted(range(len(self.speaker_path_list)), - key=lambda i: str(self.speaker_path_list[i])) - order = order[shard_idx::shard_num] - self.speaker_path_list = [self.speaker_path_list[i] for i in order] - self.listener_path_list = [self.listener_path_list[i] for i in order] - self.gt_path_list = [self.gt_path_list[i] for i in order] - print(f"[EVAL SHARD] shard {shard_idx}/{shard_num}: " - f"{len(self.speaker_path_list)} samples") - - mean_face_path = os.path.join(hydra.utils.get_original_cwd(), 'external/FaceVerse/mean_face.npy') - self.mean_face = torch.FloatTensor( - np.load(mean_face_path).astype(np.float32)).view(1, 1, -1) - std_face_path = os.path.join(hydra.utils.get_original_cwd(), 'external/FaceVerse/std_face.npy') - self.std_face = torch.FloatTensor( - np.load(std_face_path).astype(np.float32)).view(1, 1, -1) - - self._frrea_video_stride = frrea_video_stride - self.normalize_3dmm = normalize_3dmm - if normalize_3dmm == 'standard': - self._transform_3dmm = transforms.Lambda(lambda e: (e - self.mean_face) / self.std_face) - elif normalize_3dmm == 'zero_center': - self._transform_3dmm = transforms.Lambda(lambda e: e - self.mean_face) - else: - raise ValueError(f"Unknown normalize_3dmm: {normalize_3dmm}") - - self._transform = Transform(target_size, crop_size) - self._len = len(self.speaker_path_list) - - @staticmethod - def _is_archive_metadata(path): - parts = Path(path).parts - name = Path(path).name - return name.startswith('._') or name == '.DS_Store' or 'PaxHeader' in parts - - def _read_eeg_text(self, rel_path): - eeg_path = Path(self._eeg_dir) / rel_path.with_suffix('.csv') - if eeg_path.exists() and not self._is_archive_metadata(eeg_path): - return eeg_path.read_text(encoding='utf-8-sig') - - if not self.eeg_use_tar_fallback or not os.path.exists(self._eeg_tar_path): - return None - - if self._eeg_tar_members is None: - with tarfile.open(self._eeg_tar_path, 'r:gz') as tar: - self._eeg_tar_members = { - name for name in tar.getnames() - if not self._is_archive_metadata(name) - } - - member_candidates = [ - os.fspath(Path(Path(self._eeg_dir).name) / rel_path.with_suffix('.csv')).replace('\\', '/'), - os.fspath(rel_path.with_suffix('.csv')).replace('\\', '/'), - ] - member = next((name for name in member_candidates if name in self._eeg_tar_members), None) - if member is None: - return None - - with tarfile.open(self._eeg_tar_path, 'r:gz') as tar: - extracted = tar.extractfile(member) - if extracted is None: - return None - return extracted.read().decode('utf-8-sig') - - def _load_eeg(self, rel_path, total_length): - eeg_dim = len(self.eeg_target_cols) - empty_target = torch.zeros(size=(total_length, eeg_dim), dtype=torch.float32) - empty_mask = torch.zeros(size=(total_length, eeg_dim), dtype=torch.float32) - - text = self._read_eeg_text(rel_path) - if text is None: - return empty_target, empty_mask - - rows = list(csv.DictReader(io.StringIO(text))) - if len(rows) == 0: - return empty_target, empty_mask - - values = np.zeros((len(rows), eeg_dim), dtype=np.float32) - mask = np.zeros((len(rows), eeg_dim), dtype=np.float32) - for row_idx, row in enumerate(rows): - for col_idx, col in enumerate(self.eeg_target_cols): - raw_value = row.get(col, '') - if raw_value == '': - continue - try: - value = float(raw_value) - except ValueError: - continue - if not np.isfinite(value): - continue - if col in EEG_RAW_CHANNELS: - value = value / self.eeg_channel_scale - values[row_idx, col_idx] = value - mask[row_idx, col_idx] = 1.0 - - frame_to_eeg = np.floor(np.arange(total_length) / self._fps).astype(np.int64) - frame_to_eeg = np.clip(frame_to_eeg, 0, len(rows) - 1) - return ( - torch.from_numpy(values[frame_to_eeg]), - torch.from_numpy(mask[frame_to_eeg]), - ) - - def _pad_frame_clip(self, clip, target_length): - clip = clip[:target_length] - if len(clip) >= target_length: - return clip - pad_shape = (target_length - len(clip), *clip.shape[1:]) - return torch.cat((clip, clip.new_zeros(pad_shape)), dim=0) - - def __getitem__(self, index): - speaker_path = self.speaker_path_list[index] # e.g., speaker/session*/Camera-2024-06-21-103121-103102 - listener_path = self.listener_path_list[index] # e.g., listener/session*/Camera-2024-06-21-103121-103102 - - video_path = os.fspath(Path(self._video_dir) / speaker_path.with_suffix('.mp4')) - vr = VideoReader(video_path, ctx=cpu(0)) - total_length = len(vr) # length from 58 to more than 20000 - if not self.load_video_s: - del vr - - if self._split == "test": - k_select = 9 - gt_paths = self.gt_path_list[index] - listener_paths = [] - listener_paths.append(listener_path) - c = gt_paths if len(gt_paths) <= 1 else [p for p in gt_paths if p != listener_path] - listener_paths.extend(random.sample(c, k_select) \ - if len(c) >= k_select else random.choices(c, k=k_select)) - self._clip_length = total_length - - cp = random.randint(0, total_length - self._clip_length) \ - if total_length > self._clip_length else 0 - - # ========== Load Speaker Data ========== - # speaker's face video clip - speaker_video_clip = torch.zeros(size=(0,)) - if self.load_video_s: - clip = [] - for i in range(cp, cp + self._clip_length): - if i >= len(vr): - break - frame = vr[i] - img = Image.fromarray(frame.asnumpy()) - img = self._transform(img) - clip.append(img) - del vr - speaker_video_clip = torch.stack(clip, dim=0) - if total_length < self._clip_length: - speaker_video_clip = torch.cat((speaker_video_clip, - torch.zeros(size=(self._clip_length - total_length, - *speaker_video_clip.shape[1:]))), dim=0) - - # speaker's facial attribute (emotion) clip - speaker_emotion_clip = torch.zeros(size=(0,)) - if self.load_emotion_s: - emotion_path = os.fspath(Path(self._emotion_dir) / speaker_path.with_suffix('.npy')) - emotion = np.load(emotion_path) - speaker_emotion_clip = torch.from_numpy(emotion) - speaker_emotion_clip = speaker_emotion_clip[cp: cp + self._clip_length] # (25-d) - - # speaker's 3DMM coefficient (facial motion) clip - speaker_params_clip = torch.zeros(size=(0,)) - if self.load_3dmm_s: - params_path = os.fspath(Path(self._3dmm_dir) / speaker_path.with_suffix('.npy')) - params = torch.FloatTensor(np.load(params_path)).squeeze() - params = params[cp: cp + self._clip_length] - speaker_params_clip = self._transform_3dmm(params)[0] # (58-d) - - # speaker's audio feature clip - speaker_audio_clip = torch.zeros(size=(0,)) - if self.load_audio: - if self.audio_feature_type == 'wav2vec': - audio_path = os.fspath(Path(self._audio_dir) / speaker_path.with_suffix('.npy')) - speaker_audio_clip = np.load(audio_path) # (768-d) - elif self.audio_feature_type == 'mfcc': - audio_path = os.fspath(Path(self._audio_dir) / speaker_path.with_suffix('.wav')) - speaker_audio_clip = extract_audio_features(audio_path, self._fps, total_length) # (78-d) - else: - raise ValueError(f"Unknown audio feature type: {self.audio_feature_type}") - speaker_audio_clip = torch.from_numpy(speaker_audio_clip)[cp:cp + self._clip_length] - - if total_length < self._clip_length: - speaker_audio_clip = torch.cat((speaker_audio_clip, - torch.zeros(size=(self._clip_length - total_length, - speaker_audio_clip.shape[-1]))), dim=0) \ - if self.load_audio else speaker_audio_clip - speaker_emotion_clip = torch.cat((speaker_emotion_clip, - torch.zeros(size=(self._clip_length - total_length, - speaker_emotion_clip.shape[-1]))), dim=0) \ - if self.load_emotion_s else speaker_emotion_clip - speaker_params_clip = torch.cat((speaker_params_clip, - torch.zeros(size=(self._clip_length - total_length, - speaker_params_clip.shape[-1]))), dim=0) \ - if self.load_3dmm_s else speaker_params_clip - - # ========== Load Listener Data ========== - if self._split == "test": - # listener's (ground-truth) face video clip - listener_clip_length = [] - listener_video_clip = torch.zeros(size=(0,)) - if self.load_video_l: - for k, listener_path in enumerate(listener_paths): - if k != 0: - continue - - video_path = os.fspath(Path(self._video_dir) / listener_path.with_suffix('.mp4')) - vr = VideoReader(video_path, ctx=cpu(0)) - if self._frrea_video_stride and self._frrea_video_stride > 0: - # FRRea/FID only needs a subsample of listener frames (+frame 0 as the - # render reference); decode just those instead of the whole video. - indices = list(range(0, len(vr), self._frrea_video_stride)) - frames = vr.get_batch(indices).asnumpy() - clip = [self._transform(Image.fromarray(frames[j])) for j in range(frames.shape[0])] - else: - clip = [] - for f in range(len(vr)): - frame = vr[f] - img = Image.fromarray(frame.asnumpy()) - img = self._transform(img) - clip.append(img) - del vr - listener_video_clip = [torch.stack(clip, dim=0)] - listener_video_clip = (listener_video_clip + - [listener_video_clip[:1]] * (len(listener_paths) - 1)) # TODO to be modified - - listener_emotion_clip = torch.zeros(size=(0,)) - # listener's emotion ground-truths - if self.load_emotion_l: - listener_emotion_clips = [] - for listener_path in listener_paths: - emotion_path = os.fspath(Path(self._emotion_dir) / listener_path.with_suffix('.npy')) - emotion = np.load(emotion_path) - listener_clip_length.append(emotion.shape[0]) - listener_emotion_clip = torch.from_numpy(emotion) - listener_emotion_clips.append(listener_emotion_clip) - listener_emotion_clip = listener_emotion_clips - - listener_params_clip = torch.zeros(size=(0,)) - # listener's 3DMM coefficients ground-truths - if self.load_3dmm_l: - listener_params_clips = [] - for listener_path in listener_paths: - params_path = os.fspath(Path(self._3dmm_dir) / listener_path.with_suffix('.npy')) - params = torch.FloatTensor(np.load(params_path)).squeeze() - listener_params_clip = self._transform_3dmm(params)[0] - listener_params_clips.append(listener_params_clip) - listener_params_clip = listener_params_clips - - speaker_clip_length = total_length - listener_clip_length = torch.tensor(listener_clip_length) - if self.load_eeg_l: - listener_eeg_clip, listener_eeg_mask = self._load_eeg(listener_paths[0], total_length) - else: - # listener's (ground-truth) face video clip - listener_video_clip = torch.zeros(size=(0,)) - if self.load_video_l: - video_path = os.fspath(Path(self._video_dir) / listener_path.with_suffix('.mp4')) - vr = VideoReader(video_path, ctx=cpu(0)) - - clip = [] - for i in range(cp, cp + self._clip_length): - if i >= len(vr): - break - frame = vr[i] - img = Image.fromarray(frame.asnumpy()) - img = self._transform(img) - clip.append(img) - del vr - listener_video_clip = torch.stack(clip, dim=0) - - _clip_length = len(listener_video_clip) - if self.load_video_s: - listener_video_clip = torch.cat( - (listener_video_clip, speaker_video_clip[(_clip_length - self._clip_length):]), dim=0) \ - if _clip_length < self._clip_length else listener_video_clip[:self._clip_length] - else: - listener_video_clip = torch.cat( - (listener_video_clip, - torch.zeros(size=(self._clip_length - _clip_length, *listener_video_clip.shape[1:]))), dim=0) \ - if _clip_length < self._clip_length else listener_video_clip[:self._clip_length] - else: - listener_video_clip = torch.zeros(size=(self._clip_length, )) - - # listener's (ground-truth) facial attribute (emotion) clip - listener_emotion_clip = torch.zeros(size=(0,)) - if self.load_emotion_l: - emotion_path = os.fspath(Path(self._emotion_dir) / listener_path.with_suffix('.npy')) - emotion = np.load(emotion_path) - assert self.load_emotion_s, "Loading speaker's emotion is required for listener's emotion at the moment" - listener_emotion_clip = torch.from_numpy(emotion)[cp: cp + self._clip_length] - _clip_length = len(listener_emotion_clip) - listener_emotion_clip = torch.cat( - (listener_emotion_clip, speaker_emotion_clip[(_clip_length - self._clip_length):]), dim=0) \ - if _clip_length < self._clip_length else listener_emotion_clip[:self._clip_length] - - # speaker's (ground-truth) 3DMM coefficient (facial motion) clip - listener_params_clip = torch.zeros(size=(0,)) - if self.load_3dmm_l: - params_path = os.fspath(Path(self._3dmm_dir) / listener_path.with_suffix('.npy')) - params = torch.FloatTensor(np.load(params_path)).squeeze() - assert self.load_3dmm_s, "Loading speaker's 3dmm is required for listener's 3dmm at the moment" - params = params[cp: cp + self._clip_length] - listener_params_clip = self._transform_3dmm(params)[0] - _clip_length = len(listener_params_clip) - listener_params_clip = torch.cat( - (listener_params_clip, speaker_params_clip[(_clip_length - self._clip_length):]), dim=0) \ - if _clip_length < self._clip_length else listener_params_clip[:self._clip_length] - - actual_clip_len = min(total_length, self._clip_length) if self._clip_length is not None else total_length - speaker_clip_length = torch.tensor(actual_clip_len) - listener_clip_length = torch.tensor(actual_clip_len) - if self.load_eeg_l: - listener_eeg_clip, listener_eeg_mask = self._load_eeg(listener_path, total_length) - listener_eeg_clip = self._pad_frame_clip(listener_eeg_clip[cp: cp + self._clip_length], - self._clip_length) - listener_eeg_mask = self._pad_frame_clip(listener_eeg_mask[cp: cp + self._clip_length], - self._clip_length) - - sample = ( - speaker_audio_clip, - speaker_video_clip, - speaker_emotion_clip, - speaker_params_clip, - listener_video_clip, - listener_emotion_clip, - listener_params_clip, - speaker_clip_length, - listener_clip_length, - ) - if self.load_eeg_l: - sample = sample + (listener_eeg_clip, listener_eeg_mask) - return sample - - def __len__(self): - return self._len - - -class ReactionDataloader: - def __init__(self, - train_dataset: DictConfig = None, - validation_dataset: DictConfig = None, - test_dataset: DictConfig = None, - **kwargs): - - self.seed = kwargs.pop('seed') - self.train_set_cfg = train_dataset - self.val_set_cfg = validation_dataset - self.test_set_cfg = test_dataset - self.clip_length = train_dataset.clip_length - - self.collate_fn_dict = {'none': None, - 'custom': custom_collate,} - - def get_dataloader(self, stage, collate_fn: str = 'custom'): - def worker_init_fn(worker_id): - seed = self.seed + worker_id - random.seed(seed) - np.random.seed(seed) - # torch.manual_seed(seed) - - if stage == 'fit': - train_dataset = instantiate(self.train_set_cfg) - train_loader = DataLoader(dataset=train_dataset, - collate_fn=self.collate_fn_dict[collate_fn], - batch_size=self.train_set_cfg.batch_size, - shuffle=self.train_set_cfg.shuffle, - num_workers=self.train_set_cfg.num_workers, - worker_init_fn=worker_init_fn) - - val_dataset = instantiate(self.val_set_cfg) - val_loader = DataLoader(dataset=val_dataset, - collate_fn=self.collate_fn_dict[collate_fn], - batch_size=self.val_set_cfg.batch_size, - shuffle=self.val_set_cfg.shuffle, - num_workers=self.val_set_cfg.num_workers, - worker_init_fn=worker_init_fn) - return train_loader, val_loader - - else: - test_dataset = instantiate(self.test_set_cfg) - test_loader = DataLoader(dataset=test_dataset, - collate_fn=self.collate_fn_dict[collate_fn], - batch_size=self.test_set_cfg.batch_size, - shuffle=self.test_set_cfg.shuffle, - num_workers=self.test_set_cfg.num_workers, - worker_init_fn=worker_init_fn) - return test_loader diff --git a/personalised/code/dataset/tools/util.py b/personalised/code/dataset/tools/util.py deleted file mode 100644 index 7594b86e1205110152f6fbbf3dc0a234d92e5b17..0000000000000000000000000000000000000000 --- a/personalised/code/dataset/tools/util.py +++ /dev/null @@ -1,70 +0,0 @@ -import subprocess -import soundfile as sf -import torch -import numpy as np -import torchaudio -from torchvision import transforms - - -class Transform(object): - def __init__(self, target_size=224, crop_size=224): - self.target_size = target_size - self.crop_size = crop_size - - def __call__(self, img): - normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) - - transform = transforms.Compose([ - transforms.Resize(self.target_size), - transforms.CenterCrop(self.crop_size), - transforms.ToTensor(), - normalize - ]) - - img = transform(img) - return img - - -def resample_audio(input_audio_file: str, output_audio_file: str, sample_rate: int): - p = subprocess.Popen([ - "ffmpeg", "-y", "-v", "error", "-i", input_audio_file, "-ar", str(sample_rate), output_audio_file - ]) - ret = p.wait() - assert ret == 0, "Resample audio failed!" - return output_audio_file - - -def extract_audio_features(audio_path, fps, n_frames): - audio, sr = sf.read(audio_path) - - if audio.ndim == 2: - audio = audio.mean(-1) - - frame_n_samples = int(sr / fps) - curr_length = len(audio) - - target_length = frame_n_samples * n_frames - - if curr_length > target_length: - audio = audio[:target_length] - elif curr_length < target_length: - audio = np.pad(audio, [0, target_length - curr_length]) - - shifted_n_samples = 0 - curr_feats = [] - - for i in range(n_frames): - curr_samples = audio[i * frame_n_samples:shifted_n_samples + i * frame_n_samples + frame_n_samples] - curr_mfcc = torchaudio.compliance.kaldi.mfcc(torch.from_numpy(curr_samples).float().view(1, -1), - sample_frequency=sr, use_energy=True) - curr_mfcc = curr_mfcc.transpose(0, 1) # (freq, time) - curr_mfcc_d = torchaudio.functional.compute_deltas(curr_mfcc) - curr_mfcc_dd = torchaudio.functional.compute_deltas(curr_mfcc_d) - curr_mfccs = np.stack((curr_mfcc.numpy(), curr_mfcc_d.numpy(), curr_mfcc_dd.numpy())).reshape(-1) - curr_feat = curr_mfccs - - curr_feats.append(curr_feat) - - curr_feats = np.stack(curr_feats, axis=0) - - return curr_feats \ No newline at end of file diff --git a/personalised/code/external/FaceVerse/FaceVerseModel.py b/personalised/code/external/FaceVerse/FaceVerseModel.py deleted file mode 100644 index 5345641eeccf3966467f44896f87f15882eba173..0000000000000000000000000000000000000000 --- a/personalised/code/external/FaceVerse/FaceVerseModel.py +++ /dev/null @@ -1,321 +0,0 @@ -import torch -from torch import nn -import numpy as np -import math - -from .ModelRenderer import ModelRenderer -from pytorch3d.structures import Meshes -from pytorch3d.renderer import TexturesVertex - -class FaceVerseModel(nn.Module): - def __init__(self, model_dict, batch_size=1, - focal=1315, img_size=256, use_simplification=False, device='cuda:0'): - super(FaceVerseModel, self).__init__() - - self.focal = focal - self.batch_size = batch_size - self.img_size = img_size - self.device = torch.device(device) - - self.p_mat = self._get_p_mat(device) - self.reverse_z = self._get_reverse_z(device) - self.camera_pos = self._get_camera_pose(device) - self.rotXYZ = torch.eye(3).view(1, 3, 3).repeat(3, 1, 1).view(3, 1, 3, 3).to(self.device) - - self.renderer = ModelRenderer(self.focal, self.img_size, self.device) - - if use_simplification: - self.select_id = model_dict['select_id'] - self.select_id_tris = np.vstack((self.select_id * 3, self.select_id * 3 + 1, self.select_id * 3 + 2)).transpose().flatten() - self.skinmask = torch.tensor(model_dict['skinmask_select'], requires_grad=False, device=self.device) - - self.kp_inds = torch.tensor(model_dict['keypoints_select'].reshape(-1, 1), requires_grad=False).squeeze().long().to(self.device) - - self.meanshape = torch.tensor(model_dict['meanshape'].reshape(1, -1)[:, self.select_id_tris], dtype=torch.float32, requires_grad=False, device=self.device) - self.meantex = torch.tensor(model_dict['meantex'].reshape(1, -1)[:, self.select_id_tris], dtype=torch.float32, requires_grad=False, device=self.device) - - self.idBase = torch.tensor(model_dict['idBase'][self.select_id_tris], dtype=torch.float32, requires_grad=False, device=self.device) - self.expBase = torch.tensor(model_dict['exBase'][self.select_id_tris], dtype=torch.float32, requires_grad=False, device=self.device) - self.texBase = torch.tensor(model_dict['texBase'][self.select_id_tris], dtype=torch.float32, requires_grad=False, device=self.device) - - self.tri = torch.tensor(model_dict['tri_select'], dtype=torch.int64, requires_grad=False, device=self.device) - self.point_buf = torch.tensor(model_dict['point_buf_select'], dtype=torch.int64, requires_grad=False, device=self.device) - - else: - self.skinmask = torch.tensor(model_dict['skinmask'], requires_grad=False, device=self.device) - - self.kp_inds = torch.tensor(model_dict['keypoints'].reshape(-1, 1), requires_grad=False).squeeze().long().to(self.device) - - self.meanshape = torch.tensor(model_dict['meanshape'].reshape(1, -1), dtype=torch.float32, requires_grad=False, device=self.device) - self.meantex = torch.tensor(model_dict['meantex'].reshape(1, -1), dtype=torch.float32, requires_grad=False, device=self.device) - - self.idBase = torch.tensor(model_dict['idBase'], dtype=torch.float32, requires_grad=False, device=self.device) - self.expBase = torch.tensor(model_dict['exBase'], dtype=torch.float32, requires_grad=False, device=self.device) - self.texBase = torch.tensor(model_dict['texBase'], dtype=torch.float32, requires_grad=False, device=self.device) - - self.tri = torch.tensor(model_dict['tri'], dtype=torch.int64, requires_grad=False, device=self.device) - self.point_buf = torch.tensor(model_dict['point_buf'], dtype=torch.int64, requires_grad=False, device=self.device) - - self.num_vertex = self.meanshape.shape[1] // 3 - self.id_dims = self.idBase.shape[1] - self.tex_dims = self.texBase.shape[1] - self.exp_dims = self.expBase.shape[1] - self.all_dims = self.id_dims + self.tex_dims + self.exp_dims - - self.init_coeff_tensors() - - # for tracking by landmarks - self.kp_inds_view = torch.cat([self.kp_inds[:, None] * 3, self.kp_inds[:, None] * 3 + 1, self.kp_inds[:, None] * 3 + 2], dim=1).flatten() - self.idBase_view = self.idBase[self.kp_inds_view, :].detach().clone() - self.expBase_view = self.expBase[self.kp_inds_view, :].detach().clone() - self.meanshape_view = self.meanshape[:, self.kp_inds_view].detach().clone() - - def init_coeff_tensors(self): - self.id_tensor = torch.zeros( - (self.batch_size, self.id_dims), dtype=torch.float32, - requires_grad=True, device=self.device) - - self.tex_tensor = torch.zeros( - (self.batch_size, self.tex_dims), dtype=torch.float32, - requires_grad=True, device=self.device) - - self.exp_tensor = torch.zeros( - (self.batch_size, self.exp_dims), dtype=torch.float32, - requires_grad=True, device=self.device) - - self.gamma_tensor = torch.zeros( - (self.batch_size, 27), dtype=torch.float32, - requires_grad=True, device=self.device) - - self.trans_tensor = torch.zeros( - (self.batch_size, 3), dtype=torch.float32, - requires_grad=False, device=self.device) - self.trans_tensor[:, 2] += 6 - self.trans_tensor.requires_grad = True - - self.rot_tensor = torch.zeros( - (self.batch_size, 3), dtype=torch.float32, - requires_grad=False, device=self.device) - self.rot_tensor[:, 0] += math.pi - self.rot_tensor.requires_grad = True - - def get_lms(self, vs): - lms = vs[:, self.kp_inds, :] - return lms - - def split_coeffs(self, coeffs): - id_coeff = coeffs[:, :self.id_dims] # identity(shape) coeff - exp_coeff = coeffs[:, self.id_dims:self.id_dims+self.exp_dims] # expression coeff - tex_coeff = coeffs[:, self.id_dims+self.exp_dims:self.all_dims] # texture(albedo) coeff - angles = coeffs[:, self.all_dims:self.all_dims+3] # ruler angles(x,y,z) for rotation of dim 3 - gamma = coeffs[:, self.all_dims+3:self.all_dims+30] # lighting coeff for 3 channel SH function of dim 27 - translation = coeffs[:, self.all_dims+30:] # translation coeff of dim 3 - - return id_coeff, exp_coeff, tex_coeff, angles, gamma, translation - - def merge_coeffs(self, id_coeff, exp_coeff, tex_coeff, angles, gamma, translation): - coeffs = torch.cat([id_coeff, exp_coeff, tex_coeff, angles, gamma, translation], dim=1) - return coeffs - - def merge_coeffs2(self, exp_coeff, angles, translation): - coeffs = torch.cat([exp_coeff, angles, translation], dim=1) - return coeffs - - def get_packed_tensors(self): - return self.merge_coeffs(self.id_tensor, - self.exp_tensor, - self.tex_tensor, - self.rot_tensor, self.gamma_tensor, - self.trans_tensor) - - def get_packed_tensors2(self): - return self.merge_coeffs2( - self.exp_tensor, - self.rot_tensor, - self.trans_tensor) - - def forward(self, coeffs, render=True, texture=True): - id_coeff, exp_coeff, tex_coeff, angles, gamma, translation = self.split_coeffs(coeffs) - rotation = self.compute_rotation_matrix(angles) - - if render: - vs = self.get_vs(id_coeff, exp_coeff) - vs_t = self.rigid_transform(vs, rotation, translation) - - lms_t = self.get_lms(vs_t) - lms_proj = self.project_vs(lms_t) - lms_proj = torch.stack( - [lms_proj[:, :, 0], self.img_size-lms_proj[:, :, 1]], dim=2) - face_texture = self.get_color(tex_coeff) - face_norm = self.compute_norm(vs, self.tri, self.point_buf) - face_norm_r = face_norm.bmm(rotation) - face_color = self.add_illumination(face_texture, face_norm_r, gamma) - - if texture: - face_color_tv = TexturesVertex(face_color) - mesh = Meshes(vs_t, self.tri.repeat(self.batch_size, 1, 1), face_color_tv) - rendered_img = self.renderer.alb_renderer(mesh) - else: - face_color_tv = TexturesVertex(face_color * 0 + 200) - mesh = Meshes(vs_t, self.tri.repeat(self.batch_size, 1, 1), face_color_tv) - rendered_img = self.renderer.sha_renderer(mesh) - - return {'rendered_img': rendered_img, - 'lms_proj': lms_proj, - 'face_texture': face_texture, - 'vs': vs_t, - 'tri': self.tri, # triangles - 'color': face_color} - else: - lms = self.get_vs_lms(id_coeff, exp_coeff) - lms_t = self.rigid_transform( - lms, rotation, translation) - - lms_proj = self.project_vs(lms_t) - lms_proj = torch.stack( - [lms_proj[:, :, 0], self.img_size-lms_proj[:, :, 1]], dim=2) - return {'lms_proj': lms_proj} - - def get_vs(self, id_coeff, exp_coeff): - face_shape = torch.einsum('ij,aj->ai', self.idBase, id_coeff) + \ - torch.einsum('ij,aj->ai', self.expBase, exp_coeff) + self.meanshape - face_shape = face_shape.view(self.batch_size, -1, 3) - return face_shape - - def get_vs_lms(self, id_coeff, exp_coeff): - face_shape = torch.einsum('ij,aj->ai', self.idBase_view, id_coeff) + \ - torch.einsum('ij,aj->ai', self.expBase_view, torch.abs(exp_coeff)) + self.meanshape_view - face_shape = face_shape.view(self.batch_size, -1, 3) - return face_shape - - def get_color(self, tex_coeff): - face_texture = torch.einsum('ij,aj->ai', self.texBase, tex_coeff) + self.meantex - face_texture = face_texture.view(self.batch_size, -1, 3) - return face_texture - - def get_skinmask(self): - return self.skinmask - - def _get_camera_pose(self, device): - camera_pos = torch.tensor([0.0, 0.0, 10.0], device=device).reshape(1, 1, 3) - return camera_pos - - def _get_p_mat(self, device): - half_image_width = self.img_size // 2 - p_matrix = np.array([self.focal, 0.0, half_image_width, - 0.0, self.focal, half_image_width, - 0.0, 0.0, 1.0], dtype=np.float32).reshape(1, 3, 3) - return torch.tensor(p_matrix, device=device) - - def _get_reverse_z(self, device): - reverse_z = np.reshape(np.array([1.0, 0, 0, 0, 1, 0, 0, 0, -1.0], dtype=np.float32), [1, 3, 3]) - return torch.tensor(reverse_z, device=device) - - def compute_norm(self, vs, tri, point_buf): - face_id = tri - point_id = point_buf - v1 = vs[:, face_id[:, 0], :] - v2 = vs[:, face_id[:, 1], :] - v3 = vs[:, face_id[:, 2], :] - e1 = v1 - v2 - e2 = v2 - v3 - face_norm = e1.cross(e2) - - v_norm = face_norm[:, point_id, :].sum(2) - v_norm = v_norm / (v_norm.norm(dim=2).unsqueeze(2) + 1e-9) - - return v_norm - - def project_vs(self, vs): - vs = torch.matmul(vs, self.reverse_z.repeat((self.batch_size, 1, 1))) + self.camera_pos - aug_projection = torch.matmul(vs, self.p_mat.repeat((self.batch_size, 1, 1)).permute((0, 2, 1))) - face_projection = aug_projection[:, :, :2] / torch.reshape(aug_projection[:, :, 2], [self.batch_size, -1, 1]) - return face_projection - - def compute_rotation_matrix(self, angles): - sinx = torch.sin(angles[:, 0]) - siny = torch.sin(angles[:, 1]) - sinz = torch.sin(angles[:, 2]) - cosx = torch.cos(angles[:, 0]) - cosy = torch.cos(angles[:, 1]) - cosz = torch.cos(angles[:, 2]) - - if self.batch_size != 1: - rotXYZ = self.rotXYZ.repeat(1, self.batch_size, 1, 1) - else: - rotXYZ = self.rotXYZ.detach().clone() - - rotXYZ[0, :, 1, 1] = cosx - rotXYZ[0, :, 1, 2] = -sinx - rotXYZ[0, :, 2, 1] = sinx - rotXYZ[0, :, 2, 2] = cosx - rotXYZ[1, :, 0, 0] = cosy - rotXYZ[1, :, 0, 2] = siny - rotXYZ[1, :, 2, 0] = -siny - rotXYZ[1, :, 2, 2] = cosy - rotXYZ[2, :, 0, 0] = cosz - rotXYZ[2, :, 0, 1] = -sinz - rotXYZ[2, :, 1, 0] = sinz - rotXYZ[2, :, 1, 1] = cosz - - rotation = rotXYZ[2].bmm(rotXYZ[1]).bmm(rotXYZ[0]) - - return rotation.permute(0, 2, 1) - - def add_illumination(self, face_texture, norm, gamma): - gamma = gamma.view(-1, 3, 9).clone() - gamma[:, :, 0] += 0.8 - gamma = gamma.permute(0, 2, 1) - - a0 = np.pi - a1 = 2 * np.pi / np.sqrt(3.0) - a2 = 2 * np.pi / np.sqrt(8.0) - c0 = 1 / np.sqrt(4 * np.pi) - c1 = np.sqrt(3.0) / np.sqrt(4 * np.pi) - c2 = 3 * np.sqrt(5.0) / np.sqrt(12 * np.pi) - d0 = 0.5 / np.sqrt(3.0) - - norm = norm.view(-1, 3) - nx, ny, nz = norm[:, 0], norm[:, 1], norm[:, 2] - arrH = [] - - arrH.append(a0 * c0 * (nx * 0 + 1)) - arrH.append(-a1 * c1 * ny) - arrH.append(a1 * c1 * nz) - arrH.append(-a1 * c1 * nx) - arrH.append(a2 * c2 * nx * ny) - arrH.append(-a2 * c2 * ny * nz) - arrH.append(a2 * c2 * d0 * (3 * nz.pow(2) - 1)) - arrH.append(-a2 * c2 * nx * nz) - arrH.append(a2 * c2 * 0.5 * (nx.pow(2) - ny.pow(2))) - - H = torch.stack(arrH, 1) - Y = H.view(self.batch_size, face_texture.shape[1], 9) - lighting = Y.bmm(gamma) - - face_color = face_texture * lighting - return face_color - - def rigid_transform(self, vs, rot, trans): - vs_r = torch.matmul(vs, rot) - vs_t = vs_r + trans.view(-1, 1, 3) - return vs_t - - def get_rot_tensor(self): - return self.rot_tensor - - def get_trans_tensor(self): - return self.trans_tensor - - def get_exp_tensor(self): - return self.exp_tensor - - def get_tex_tensor(self): - return self.tex_tensor - - def get_id_tensor(self): - return self.id_tensor - - def get_gamma_tensor(self): - return self.gamma_tensor - diff --git a/personalised/code/external/FaceVerse/LICENSE b/personalised/code/external/FaceVerse/LICENSE deleted file mode 100644 index ba7be5c07c029703156ea22a68e9d1a62dc796bc..0000000000000000000000000000000000000000 --- a/personalised/code/external/FaceVerse/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2022 Evelyn - -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. diff --git a/personalised/code/external/FaceVerse/ModelRenderer.py b/personalised/code/external/FaceVerse/ModelRenderer.py deleted file mode 100644 index c56aa8337fd9a100a2d6c47857e61e4cf7d7d98e..0000000000000000000000000000000000000000 --- a/personalised/code/external/FaceVerse/ModelRenderer.py +++ /dev/null @@ -1,60 +0,0 @@ -import numpy as np -import torch - -from pytorch3d.structures import Meshes -from pytorch3d.renderer import ( - look_at_view_transform, - FoVPerspectiveCameras, - PointLights, - RasterizationSettings, - MeshRenderer, - MeshRasterizer, - HardFlatShader, - TexturesVertex, - blending -) - -class ModelRenderer: - def __init__(self, focal=1315, img_size=224, device='cuda:0'): - self.img_size = img_size - self.focal = focal - self.device = device - - self.alb_renderer = self._get_renderer(albedo=True) - self.sha_renderer = self._get_renderer(albedo=False) - - def _get_renderer(self, albedo=True): - R, T = look_at_view_transform(10, 0, 0) # camera's position - cameras = FoVPerspectiveCameras(device=self.device, R=R, T=T, znear=0.01, zfar=50, - fov=2 * np.arctan(self.img_size // 2 / self.focal) * 180. / np.pi) - - if albedo: - lights = PointLights(device=self.device, location=[[0.0, 0.0, 1e5]], - ambient_color=[[1, 1, 1]], - specular_color=[[0., 0., 0.]], diffuse_color=[[0., 0., 0.]]) - else: - lights = PointLights(device=self.device, location=[[0.0, 0.0, 1e5]], - ambient_color=[[0.1, 0.1, 0.1]], - specular_color=[[0.0, 0.0, 0.0]], diffuse_color=[[0.95, 0.95, 0.95]]) - - raster_settings = RasterizationSettings( - image_size=self.img_size, - blur_radius=0.0, - faces_per_pixel=1, - ) - blend_params = blending.BlendParams(background_color=[0, 0, 0]) - - renderer = MeshRenderer( - rasterizer=MeshRasterizer( - cameras=cameras, - raster_settings=raster_settings - ), - shader=HardFlatShader( - device=self.device, - cameras=cameras, - lights=lights, - blend_params=blend_params - ) - ) - return renderer - diff --git a/personalised/code/external/FaceVerse/__init__.py b/personalised/code/external/FaceVerse/__init__.py deleted file mode 100644 index 38ba6ca7563a27fc9dc1948b04c4248c9cbfee88..0000000000000000000000000000000000000000 --- a/personalised/code/external/FaceVerse/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from .FaceVerseModel import FaceVerseModel -import numpy as np - -def get_faceverse(**kargs): - # Use the caller-provided absolute path when available; the hardcoded relative path - # breaks once Hydra changes the working directory to the run output dir. - path = kargs.pop('path', 'external/FaceVerse/data/faceverse_simple_v2.npy') - faceverse_dict = np.load(path, allow_pickle=True).item() - faceverse_model = FaceVerseModel(faceverse_dict, **kargs) - return faceverse_model, faceverse_dict - - - diff --git a/personalised/code/external/FaceVerse/data/faceverse_simple_v2.npy b/personalised/code/external/FaceVerse/data/faceverse_simple_v2.npy deleted file mode 100644 index 53726c387cad7910c858bda21304090d4c6d9899..0000000000000000000000000000000000000000 --- a/personalised/code/external/FaceVerse/data/faceverse_simple_v2.npy +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:56c574bfd44afdd3d3197f3da54c14d3ab52ae9499caa76748c666a798ebb44a -size 160199539 diff --git a/personalised/code/external/FaceVerse/mean_face.npy b/personalised/code/external/FaceVerse/mean_face.npy deleted file mode 100644 index 5140b79c85a712a2cb60d8eb809659d5f03d2cb8..0000000000000000000000000000000000000000 --- a/personalised/code/external/FaceVerse/mean_face.npy +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9160971fce93790a241d7bb666f5db8be0dadb56d0caee7e629336114a13ec82 -size 360 diff --git a/personalised/code/external/FaceVerse/reference_full.npy b/personalised/code/external/FaceVerse/reference_full.npy deleted file mode 100644 index 1fdd7e9e77db464b019173bbc1cfeb2f36a5f0b3..0000000000000000000000000000000000000000 --- a/personalised/code/external/FaceVerse/reference_full.npy +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:406b632c4ebd0ca0408e17e0981048f963e3395a72b28f3dca1a6227b20538fa -size 2072 diff --git a/personalised/code/external/FaceVerse/std_face.npy b/personalised/code/external/FaceVerse/std_face.npy deleted file mode 100644 index 8614b369ebba91b0c87c4a54d3ee433359d8f886..0000000000000000000000000000000000000000 --- a/personalised/code/external/FaceVerse/std_face.npy +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5324c112a5cd444a07bb15dc35c6d08bcfb0546e96f6c2e1ea643c6f8fedd5c9 -size 360 diff --git a/personalised/code/framework/feature_extractor/wav2vec.py b/personalised/code/framework/feature_extractor/wav2vec.py deleted file mode 100644 index 14ff991e09e38a074a9eb730f19d88e7a9f2c3b7..0000000000000000000000000000000000000000 --- a/personalised/code/framework/feature_extractor/wav2vec.py +++ /dev/null @@ -1,216 +0,0 @@ -# Adapted from https://github.com/fudan-generative-vision/hallo/blob/main/hallo/models/wav2vec.py -# Original author: Fusion Lab: Generative Vision Lab of Fudan University -# License: MIT License - -# pylint: disable=R0901 -# src/models/wav2vec.py - -""" -This module defines the Wav2Vec model, which is a pre-trained model for speech recognition and understanding. -It inherits from the Wav2Vec2Model class in the transformers library and provides additional functionalities -such as feature extraction and encoding. - -Classes: - Wav2VecModel: Inherits from Wav2Vec2Model and adds additional methods for feature extraction and encoding. - -Functions: - linear_interpolation: Interpolates the features based on the sequence length. -""" - -import torch.nn.functional as F -from transformers import Wav2Vec2Model -from transformers.modeling_outputs import BaseModelOutput - - -class Wav2VecModel(Wav2Vec2Model): - """ - Wav2VecModel is a custom model class that extends the Wav2Vec2Model class from the transformers library. - It inherits all the functionality of the Wav2Vec2Model and adds additional methods for feature extraction and encoding. - ... - - Attributes: - base_model (Wav2Vec2Model): The base Wav2Vec2Model object. - - Methods: - forward(input_values, seq_len, attention_mask=None, mask_time_indices=None - , output_attentions=None, output_hidden_states=None, return_dict=None): - Forward pass of the Wav2VecModel. - It takes input_values, seq_len, and other optional parameters as input and returns the output of the base model. - - feature_extract(input_values, seq_len): - Extracts features from the input_values using the base model. - - encode(extract_features, attention_mask=None, mask_time_indices=None, output_attentions=None, output_hidden_states=None, return_dict=None): - Encodes the extracted features using the base model and returns the encoded features. - """ - - def forward( - self, - input_values, - seq_len, - attention_mask=None, - mask_time_indices=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - ): - """ - Forward pass of the Wav2Vec model. - - Args: - self: The instance of the model. - input_values: The input values (waveform) to the model. - seq_len: The sequence length of the input values. - attention_mask: Attention mask to be used for the model. - mask_time_indices: Mask indices to be used for the model. - output_attentions: If set to True, returns attentions. - output_hidden_states: If set to True, returns hidden states. - return_dict: If set to True, returns a BaseModelOutput instead of a tuple. - - Returns: - The output of the Wav2Vec model. - """ - self.config.output_attentions = True - - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - extract_features = self.feature_extractor(input_values) - extract_features = extract_features.transpose(1, 2) - - print(f"Linear interpolation from original length {extract_features.shape[1]} to {seq_len}") - extract_features = linear_interpolation(extract_features, seq_len=seq_len) - - if attention_mask is not None: - # compute reduced attention_mask corresponding to feature vectors - attention_mask = self._get_feature_vector_attention_mask( - extract_features.shape[1], attention_mask, add_adapter=False - ) - - hidden_states, extract_features = self.feature_projection(extract_features) - hidden_states = self._mask_hidden_states( - hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask - ) - - encoder_outputs = self.encoder( - hidden_states, - attention_mask=attention_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - # hidden_states = encoder_outputs[0] - hidden_states = encoder_outputs.last_hidden_state - - if self.adapter is not None: - hidden_states = self.adapter(hidden_states) - - if not return_dict: - return (hidden_states,) + encoder_outputs[1:] - return BaseModelOutput( - last_hidden_state=hidden_states, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - ) - - def feature_extract( - self, - input_values, - seq_len, - ): - """ - Extracts features from the input values and returns the extracted features. - - Parameters: - input_values (torch.Tensor): The input values to be processed. - seq_len (torch.Tensor): The sequence lengths of the input values. - - Returns: - extracted_features (torch.Tensor): The extracted features from the input values. - """ - extract_features = self.feature_extractor(input_values) - extract_features = extract_features.transpose(1, 2) - extract_features = linear_interpolation(extract_features, seq_len=seq_len) - - return extract_features - - def encode( - self, - extract_features, - attention_mask=None, - mask_time_indices=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - ): - """ - Encodes the input features into the output space. - - Args: - extract_features (torch.Tensor): The extracted features from the audio signal. - attention_mask (torch.Tensor, optional): Attention mask to be used for padding. - mask_time_indices (torch.Tensor, optional): Masked indices for the time dimension. - output_attentions (bool, optional): If set to True, returns the attention weights. - output_hidden_states (bool, optional): If set to True, returns all hidden states. - return_dict (bool, optional): If set to True, returns a BaseModelOutput instead of the tuple. - - Returns: - The encoded output features. - """ - self.config.output_attentions = True - - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if attention_mask is not None: - # compute reduced attention_mask corresponding to feature vectors - attention_mask = self._get_feature_vector_attention_mask( - extract_features.shape[1], attention_mask, add_adapter=False - ) - - hidden_states, extract_features = self.feature_projection(extract_features) - hidden_states = self._mask_hidden_states( - hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask - ) - - encoder_outputs = self.encoder( - hidden_states, - attention_mask=attention_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - hidden_states = encoder_outputs[0] - - if self.adapter is not None: - hidden_states = self.adapter(hidden_states) - - if not return_dict: - return (hidden_states,) + encoder_outputs[1:] - return BaseModelOutput( - last_hidden_state=hidden_states, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - ) - - -def linear_interpolation(features, seq_len): - """ - Transpose the features to interpolate linearly. - - Args: - features (torch.Tensor): The extracted features to be interpolated. - seq_len (torch.Tensor): The sequence lengths of the features. - - Returns: - torch.Tensor: The interpolated features. - """ - features = features.transpose(1, 2) - output_features = F.interpolate(features, size=seq_len, align_corners=True, mode='linear') - return output_features.transpose(1, 2) diff --git a/personalised/code/framework/g2p_delta/__init__.py b/personalised/code/framework/g2p_delta/__init__.py deleted file mode 100644 index 404ec802c7be3664b1d4a65fdedf54460627663c..0000000000000000000000000000000000000000 --- a/personalised/code/framework/g2p_delta/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .model import G2PDeltaModel - -__all__ = ["G2PDeltaModel"] diff --git a/personalised/code/framework/g2p_delta/model.py b/personalised/code/framework/g2p_delta/model.py deleted file mode 100644 index aac7e555b55569b2709cb67183be849960dec138..0000000000000000000000000000000000000000 --- a/personalised/code/framework/g2p_delta/model.py +++ /dev/null @@ -1,325 +0,0 @@ -"""Generic-to-Personal residual denoising adapter. - -The verified Generic Offline diffusion model remains frozen. A listener -condition produces a bounded, zero-initialised correction at the denoiser's -existing ``to_emotion_feat`` projection, so a fresh model is exactly the -Generic model while retaining its prior and EEG paths. -""" - -from __future__ import annotations - -from collections import OrderedDict -from typing import Optional - -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class PersonalityEncoder(nn.Module): - def __init__(self, output_dim: int, hidden_dim: int = 128, dropout: float = 0.1): - super().__init__() - self.net = nn.Sequential( - nn.LayerNorm(5), - nn.Linear(5, hidden_dim), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(hidden_dim, output_dim), - ) - - def forward(self, personality: torch.Tensor) -> torch.Tensor: - return F.normalize(self.net(personality.float()), dim=-1) - - -class HistoryEncoder(nn.Module): - """Checkpoint-free LHFB encoder for a 58-D historical 3DMM sequence.""" - - def __init__(self, output_dim: int, hidden_dim: int = 128, dropout: float = 0.1): - super().__init__() - self.temporal = nn.Sequential( - nn.Conv1d(58, hidden_dim, kernel_size=5, padding=2), - nn.GELU(), - nn.Conv1d(hidden_dim, hidden_dim, kernel_size=5, padding=4, dilation=2), - nn.GELU(), - ) - self.project = nn.Sequential( - nn.LayerNorm(hidden_dim * 2), - nn.Linear(hidden_dim * 2, output_dim), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(output_dim, output_dim), - ) - - def forward(self, history: torch.Tensor) -> torch.Tensor: - if history.dim() != 3 or history.shape[-1] != 58: - raise ValueError(f"LHFB must have shape (B,T,58), got {tuple(history.shape)}") - features = self.temporal(history.float().transpose(1, 2)) - pooled = torch.cat( - (features.mean(dim=-1), features.std(dim=-1, unbiased=False)), dim=-1 - ) - return F.normalize(self.project(pooled), dim=-1) - - -class ConditionFusion(nn.Module): - def __init__(self, embed_dim: int, dropout: float = 0.1): - super().__init__() - self.gate = nn.Sequential( - nn.LayerNorm(embed_dim * 2), - nn.Linear(embed_dim * 2, embed_dim), - nn.Sigmoid(), - ) - self.out = nn.Sequential( - nn.LayerNorm(embed_dim), - nn.Dropout(dropout), - nn.Linear(embed_dim, embed_dim), - ) - - def forward(self, history: torch.Tensor, personality: torch.Tensor) -> torch.Tensor: - gate = self.gate(torch.cat((history, personality), dim=-1)) - return F.normalize(self.out(gate * history + (1.0 - gate) * personality), dim=-1) - - -class ResidualDenoisingAdapter(nn.Module): - """Low-rank correction applied at every reverse-diffusion denoiser call.""" - - def __init__( - self, - hidden_dim: int, - condition_dim: int, - output_dim: int = 25, - rank: int = 128, - max_scale: float = 0.15, - ): - super().__init__() - self.hidden_norm = nn.LayerNorm(hidden_dim) - self.hidden_down = nn.Linear(hidden_dim, rank) - self.condition = nn.Linear(condition_dim, rank) - self.output = nn.Linear(rank, output_dim) - self.gate_logit = nn.Parameter(torch.tensor(0.0)) - self.max_scale = float(max_scale) - nn.init.zeros_(self.output.weight) - nn.init.zeros_(self.output.bias) - - def forward(self, hidden: torch.Tensor, condition: torch.Tensor) -> torch.Tensor: - # hidden: (T, B*num_preds, D); one listener condition is broadcast to samples. - batch = hidden.shape[1] - if condition.shape[0] == 1 and batch != 1: - condition = condition.expand(batch, -1) - elif condition.shape[0] != batch: - raise ValueError( - f"Condition batch {condition.shape[0]} cannot broadcast to denoiser batch {batch}" - ) - joint = self.hidden_down(self.hidden_norm(hidden)) - joint = joint + self.condition(condition).unsqueeze(0) - delta = self.output(F.gelu(joint)) - scale = self.max_scale * torch.sigmoid(self.gate_logit) - return scale * delta - - -class CoarsePriorAdapter(nn.Module): - """Small zero-init listener correction to the Generic coarse expression plan.""" - - def __init__(self, condition_dim: int, coarse_hidden: int = 256): - super().__init__() - self.net = nn.Sequential( - nn.LayerNorm(condition_dim), - nn.Linear(condition_dim, coarse_hidden), - nn.GELU(), - nn.Linear(coarse_hidden, 2 * coarse_hidden), - ) - nn.init.zeros_(self.net[-1].weight) - nn.init.zeros_(self.net[-1].bias) - - def forward(self, condition: torch.Tensor): - gamma, beta = self.net(condition).chunk(2, dim=-1) - return 0.1 * torch.tanh(gamma), 0.1 * torch.tanh(beta) - - -class G2PDeltaModel(nn.Module): - """Thin personalised extension around the final Generic Offline model.""" - - VALID_MODES = {"3dmm_only", "personality_only", "3dmm_personality"} - - def __init__(self, cfg, main_net: nn.Module): - super().__init__() - self.main_net = main_net - args = cfg.main_model.args - self.personal_condition_mode = str( - args.get("personal_condition_mode", "personality_only") - ) - if self.personal_condition_mode not in self.VALID_MODES: - raise ValueError(f"Unsupported personal condition: {self.personal_condition_mode}") - - embed_dim = int(args.get("embed_dim", 256)) - dropout = float(args.get("condition_dropout", 0.1)) - self.history_encoder = ( - HistoryEncoder(embed_dim, int(args.get("history_hidden_dim", 128)), dropout) - if self.personal_condition_mode in {"3dmm_only", "3dmm_personality"} - else None - ) - self.personality_encoder = ( - PersonalityEncoder(embed_dim, int(args.get("personality_hidden_dim", 128)), dropout) - if self.personal_condition_mode in {"personality_only", "3dmm_personality"} - else None - ) - self.condition_fusion = ( - ConditionFusion(embed_dim, dropout) - if self.personal_condition_mode == "3dmm_personality" - else None - ) - - modules = OrderedDict(self.main_net.named_modules()) - denoiser = modules.get("diffusion_decoder.model") - if denoiser is None: - raise ValueError("Generic model has no diffusion_decoder.model") - if not hasattr(denoiser, "to_emotion_feat"): - raise ValueError("Generic denoiser has no to_emotion_feat projection") - self._denoiser = denoiser - hidden_dim = int(getattr(denoiser, "latent_dim", 512)) - self.delta_adapter = ResidualDenoisingAdapter( - hidden_dim=hidden_dim, - condition_dim=embed_dim, - output_dim=25, - rank=int(args.get("delta_rank", 128)), - max_scale=float(args.get("delta_max_scale", 0.15)), - ) - self.coarse_adapter = ( - CoarsePriorAdapter(embed_dim, denoiser.coarse_gru.hidden_size) - if bool(args.get("personalize_coarse", True)) and hasattr(denoiser, "coarse_gru") - else None - ) - self.anchor_weight = float(args.get("anchor_weight", 1.0e-3)) - self._current_condition: Optional[torch.Tensor] = None - self._last_delta_energy = torch.tensor(0.0) - self._hook = denoiser.to_emotion_feat.register_forward_hook(self._delta_hook) - - # Freeze the verified Generic backbone. EEG is selectively unfrozen by the trainer. - for parameter in self.main_net.parameters(): - parameter.requires_grad = False - - # Compatibility with the inherited personalised trainer. - self.person_encoder = None - - def _delta_hook(self, _module, inputs, output): - if self._current_condition is None: - self._last_delta_energy = output.new_tensor(0.0) - return output - hidden = inputs[0] - delta = self.delta_adapter(hidden, self._current_condition) - self._last_delta_energy = delta.square().mean() - return output + delta - - def encode_person_condition(self, p=None, personality=None): - history_embedding = personality_embedding = None - if self.history_encoder is not None: - if p is None or p.numel() == 0: - raise ValueError("LHFB condition requires non-empty personal 3DMM history") - history_embedding = self.history_encoder(p) - if self.personality_encoder is not None: - if personality is None or personality.numel() == 0: - raise ValueError("Personality condition requires the listener Big-Five vector") - if personality.dim() == 1: - personality = personality.unsqueeze(0) - personality_embedding = self.personality_encoder(personality) - if self.personal_condition_mode == "3dmm_only": - return history_embedding - if self.personal_condition_mode == "personality_only": - return personality_embedding - return self.condition_fusion(history_embedding, personality_embedding) - - def set_person_condition(self, p=None, personality=None): - condition = self.encode_person_condition(p=p, personality=personality) - self._current_condition = condition - if self.coarse_adapter is not None: - self._denoiser._person_coarse_film = self.coarse_adapter(condition) - return condition - - def clear_person_condition(self): - self._current_condition = None - if self.coarse_adapter is not None: - self._denoiser._person_coarse_film = None - - def forward(self, x, p=None, personality=None): - self.set_person_condition(p=p, personality=personality) - try: - output = self.main_net(**x) - regular = self.anchor_weight * self._last_delta_energy - finally: - self.clear_person_condition() - return output, regular - - def eeg_head(self): - return getattr(self.main_net, "eeg_head", None) - - def has_eeg_head(self): - return self.eeg_head() is not None - - def set_eeg_head_requires_grad(self, requires_grad=True): - if not self.has_eeg_head(): - raise RuntimeError("Generic EEG prediction head is disabled") - for parameter in self.eeg_head().parameters(): - parameter.requires_grad = requires_grad - - def modifier_parameters(self, include_eeg_head=False): - modules = [ - self.history_encoder, - self.personality_encoder, - self.condition_fusion, - self.delta_adapter, - self.coarse_adapter, - ] - for module in modules: - if module is not None: - yield from module.parameters() - if include_eeg_head: - self.set_eeg_head_requires_grad(True) - yield from self.eeg_head().parameters() - - def modifier_state_dict(self, include_eeg_head=False): - names = ( - "history_encoder", - "personality_encoder", - "condition_fusion", - "delta_adapter", - "coarse_adapter", - ) - state = {} - for name in names: - module = getattr(self, name) - if module is not None: - state.update({f"{name}.{k}": v for k, v in module.state_dict().items()}) - if include_eeg_head and self.has_eeg_head(): - state.update({f"eeg_head.{k}": v for k, v in self.eeg_head().state_dict().items()}) - return state - - def load_modifier_state_dict(self, state_dict): - for name in ( - "history_encoder", - "personality_encoder", - "condition_fusion", - "delta_adapter", - "coarse_adapter", - ): - module = getattr(self, name) - selected = { - key[len(name) + 1 :]: value - for key, value in state_dict.items() - if key.startswith(f"{name}.") - } - if selected: - if module is None: - raise ValueError(f"Checkpoint contains {name}, but it is disabled") - module.load_state_dict(selected) - eeg_state = { - key[len("eeg_head.") :]: value - for key, value in state_dict.items() - if key.startswith("eeg_head.") - } - if eeg_state: - self.eeg_head().load_state_dict(eeg_state) - - def train(self, mode: bool = True): - super().train(mode) - # The final Generic model stays deterministic/frozen; only adapters train. - self.main_net.eval() - return self diff --git a/personalised/code/framework/metrics/ACC.py b/personalised/code/framework/metrics/ACC.py deleted file mode 100644 index 54d72e19d21c9da1848f121d8f3c53701281be79..0000000000000000000000000000000000000000 --- a/personalised/code/framework/metrics/ACC.py +++ /dev/null @@ -1,4 +0,0 @@ -from scipy.spatial.distance import pdist -import numpy as np -import torch -import os diff --git a/personalised/code/framework/metrics/FID.py b/personalised/code/framework/metrics/FID.py deleted file mode 100644 index e296198cb82b8b071f02dea918aeefeb8e59e8ea..0000000000000000000000000000000000000000 --- a/personalised/code/framework/metrics/FID.py +++ /dev/null @@ -1,140 +0,0 @@ -"""FRRea metric: Frechet Inception Distance (FID) over rendered facial reaction frames. - -Realism of generated facial reaction video clips is assessed with FID (denoted FRRea): -features are extracted from generated ("fake") and real listener frames with a -pretrained InceptionV3, then the Frechet distance between the two Gaussians is computed. - -FID is a distribution-level metric (mean + covariance over the whole frame set), so it -is NOT a per-sample average and cannot be merged across data shards by averaging. When -evaluation is sharded across GPUs, each shard should only dump its frames to a shared -directory; FID is then computed once over all frames with this module. - -Usage (standalone, over two image directories): - python -m framework.metrics.FID --fake --real -""" -import argparse -import os - -import numpy as np -import torch -import torch.nn as nn -from scipy import linalg - -try: - import cv2 -except ImportError: - cv2 = None - -from torchvision.models import inception_v3, Inception_V3_Weights - -IMG_EXTS = ('.png', '.jpg', '.jpeg', '.bmp') - - -class InceptionFeatures(nn.Module): - """InceptionV3 pool3 (2048-d) feature extractor, standard for FID.""" - - def __init__(self, device='cpu'): - super().__init__() - net = inception_v3(weights=Inception_V3_Weights.IMAGENET1K_V1, - aux_logits=True, transform_input=True) - net.fc = nn.Identity() # output the 2048-d pooled features - net.eval() - self.net = net.to(device) - self.device = device - - @torch.no_grad() - def forward(self, x): - # x: (B, 3, H, W) float in [0, 1] - x = nn.functional.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False) - feat = self.net(x) - if isinstance(feat, tuple): # eval() should disable aux, but be safe - feat = feat[0] - return feat - - -def _list_images(directory): - files = [] - for root, _, names in os.walk(directory): - for n in sorted(names): - if n.lower().endswith(IMG_EXTS): - files.append(os.path.join(root, n)) - return sorted(files) - - -def _load_batch(paths): - """Load images as RGB float tensor in [0, 1]. Frames are saved by cv2 (BGR).""" - imgs = [] - for p in paths: - bgr = cv2.imread(p, cv2.IMREAD_COLOR) - if bgr is None: - continue - rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) - imgs.append(torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0) - if not imgs: - return None - return torch.stack(imgs, dim=0) - - -@torch.no_grad() -def extract_features_from_dir(directory, extractor, batch_size=64): - paths = _list_images(directory) - if not paths: - raise ValueError(f"No images found in {directory}") - feats = [] - for i in range(0, len(paths), batch_size): - batch = _load_batch(paths[i:i + batch_size]) - if batch is None: - continue - batch = batch.to(extractor.device) - feats.append(extractor(batch).cpu().numpy()) - feats = np.concatenate(feats, axis=0) - return feats, len(paths) - - -def frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6): - diff = mu1 - mu2 - covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False) - if not np.isfinite(covmean).all(): - # add small jitter to the diagonal to make the product positive definite - offset = np.eye(sigma1.shape[0]) * eps - covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset)) - if np.iscomplexobj(covmean): - covmean = covmean.real - return float(diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * np.trace(covmean)) - - -def _stats(feats): - mu = np.mean(feats, axis=0) - sigma = np.cov(feats, rowvar=False) - return mu, sigma - - -def compute_fid_from_dirs(fake_dir, real_dir, device=None, batch_size=64): - device = device or ('cuda' if torch.cuda.is_available() else 'cpu') - extractor = InceptionFeatures(device=device) - fake_feats, n_fake = extract_features_from_dir(fake_dir, extractor, batch_size) - real_feats, n_real = extract_features_from_dir(real_dir, extractor, batch_size) - mu_f, sig_f = _stats(fake_feats) - mu_r, sig_r = _stats(real_feats) - fid = frechet_distance(mu_f, sig_f, mu_r, sig_r) - return fid, n_fake, n_real - - -def main(): - ap = argparse.ArgumentParser(description="Compute FRRea (FID) between fake and real frame dirs.") - ap.add_argument('--fake', required=True, help='directory of generated (fake) frames') - ap.add_argument('--real', required=True, help='directory of real listener frames') - ap.add_argument('--device', default=None) - ap.add_argument('--batch_size', type=int, default=64) - args = ap.parse_args() - - if cv2 is None: - raise ImportError("opencv (cv2) is required to read frames for FID.") - - fid, n_fake, n_real = compute_fid_from_dirs(args.fake, args.real, args.device, args.batch_size) - print(f"fake frames: {n_fake}, real frames: {n_real}") - print(f"FRRea (FID) = {fid:.6f}") - - -if __name__ == '__main__': - main() diff --git a/personalised/code/framework/metrics/FRC.py b/personalised/code/framework/metrics/FRC.py deleted file mode 100644 index 716d876f0828b0e1b031cb409e330eeb3a7167d9..0000000000000000000000000000000000000000 --- a/personalised/code/framework/metrics/FRC.py +++ /dev/null @@ -1,208 +0,0 @@ -import warnings -import numpy as np -from tslearn.metrics import dtw -from functools import partial -import multiprocessing as mp -from tqdm import tqdm - - -def corrcoef(x, y=None, rowvar=True, bias=np._NoValue, ddof=np._NoValue, *, - dtype=None): - if bias is not np._NoValue or ddof is not np._NoValue: - # 2015-03-15, 1.10 - warnings.warn('bias and ddof have no effect and are deprecated', - DeprecationWarning, stacklevel=3) - c = np.cov(x, y, rowvar, dtype=dtype) - - try: - d = np.diag(c) - except ValueError: - # scalar covariance - # nan if incorrect value (nan, inf, 0), 1 otherwise - return c / c - stddev = np.sqrt(d.real) - - c /= stddev[:, None] - c /= stddev[None, :] - c = np.nan_to_num(c) - - # Clip real and imaginary parts to [-1, 1]. This does not guarantee - # abs(a[i,j]) <= 1 for complex arrays, but is the best we can do without - # excessive work. - np.clip(c.real, -1, 1, out=c.real) - if np.iscomplexobj(c): - np.clip(c.imag, -1, 1, out=c.imag) - return c - - -def concordance_correlation_coefficient(y_true, y_pred, - sample_weight=None, - multioutput='uniform_average'): - """Concordance correlation coefficient. - The concordance correlation coefficient is a measure of inter-rater agreement. - It measures the deviation of the relationship between predicted and true values - from the 45 degree angle. - Read more: https://en.wikipedia.org/wiki/Concordance_correlation_coefficient - Original paper: Lawrence, I., and Kuei Lin. "A concordance correlation coefficient to evaluate reproducibility." Biometrics (1989): 255-268. - Parameters - ---------- - y_true : array-like of shape = (n_samples) or (n_samples, n_outputs) - Ground truth (correct) target values. - y_pred : array-like of shape = (n_samples) or (n_samples, n_outputs) - Estimated target values. - Returns - ------- - loss : A float in the range [-1,1]. A value of 1 indicates perfect agreement - between the true and the predicted values. - Examples - -------- - >>> from sklearn.metrics import concordance_correlation_coefficient - >>> y_true = [3, -0.5, 2, 7] - >>> y_pred = [2.5, 0.0, 2, 8] - >>> concordance_correlation_coefficient(y_true, y_pred) - 0.97678916827853024 - """ - - # y_true.shape: (seq_len, dim); y_pred.shape: (seq_len, dim) - if len(y_true.shape) > 1: - ccc_list = [] - for i in range(y_true.shape[1]): # dim==25 - cor = corrcoef(y_true[:, i], y_pred[:, i])[0][1] - mean_true = np.mean(y_true[:, i]) - - mean_pred = np.mean(y_pred[:, i]) - - var_true = np.var(y_true[:, i]) - var_pred = np.var(y_pred[:, i]) - - sd_true = np.std(y_true[:, i]) - sd_pred = np.std(y_pred[:, i]) - - numerator = 2 * cor * sd_true * sd_pred - - denominator = var_true + var_pred + (mean_true - mean_pred) ** 2 - - ccc = numerator / (denominator + 1e-8) - - ccc_list.append(ccc) - ccc = np.mean(ccc_list) - else: - cor = np.corrcoef(y_true, y_pred)[0][1] - mean_true = np.mean(y_true) - mean_pred = np.mean(y_pred) - - var_true = np.var(y_true) - var_pred = np.var(y_pred) - - sd_true = np.std(y_true) - sd_pred = np.std(y_pred) - - numerator = 2 * cor * sd_true * sd_pred - - denominator = var_true + var_pred + (mean_true - mean_pred) ** 2 - ccc = numerator / (denominator + 1e-8) - return ccc, ccc_list - - -def _func(target, pred): - # target: (10, l, dim) - # pred: (10, l, dim) - # num_preds = pred.shape[0] - max_ccc_sum = 0 - for i in range(pred.shape[0]): - ccc_list = [] - for j in range(target.shape[0]): - ccc, lst = concordance_correlation_coefficient(target[j].numpy(), pred[i].numpy()) - # if i == 0: - # print(f"25 dimension ccc lst: {lst} at j={j}") - ccc_list.append(ccc) - max_ccc_sum += max(ccc_list) - return max_ccc_sum - - -def _func_star(args): - target, pred = args - return _func(target, pred) - -def compute_FRC(preds, targets, p=1): - tasks = list(zip(targets, preds)) - FRC_list = [] - - with mp.Pool(processes=p) as pool: - for result in tqdm( - pool.imap(_func_star, tasks), - total=len(tasks), - desc="Computing FRC" - ): - FRC_list.append(result) - - return np.mean(FRC_list) - - -# def compute_FRC_mp(args, pred, em, val_test='test', p=1): -# # pred: N 10 750 dim -# # em: N 750 dim -# if val_test == 'val': -# # neighbour_matrix = np.load(os.path.join(args.dataset_path, 'neighbour_emotion_val.npy')) -# neighbour_matrix = np.load(os.path.join(args.dataset_path, 'person_specific_masked_neighbour_emotion_val.npy')) -# else: -# # neighbour_matrix = np.load(os.path.join(args.dataset_path, 'neighbour_emotion_test.npy')) -# neighbour_matrix = np.load(os.path.join(args.dataset_path, 'person_specific_masked_neighbour_emotion_test.npy')) -# -# FRC_list = [] -# with mp.Pool(processes=p) as pool: -# # use map -# _func_partial = partial(_func, em=em.numpy()) -# FRC_list += pool.starmap(_func_partial, zip(neighbour_matrix, pred.numpy())) -# return np.mean(FRC_list) - - -# def _func(k_neighbour_matrix, k_pred, em=None): -# neighbour_index = np.argwhere(k_neighbour_matrix == 1).reshape(-1) -# neighbour_index_len = len(neighbour_index) -# max_ccc_sum = 0 -# for i in range(k_pred.shape[0]): -# ccc_list = [] -# for n_index in range(neighbour_index_len): -# emotion = em[neighbour_index[n_index]] -# ccc = concordance_correlation_coefficient(emotion, k_pred[i]) -# ccc_list.append(ccc) -# max_ccc_sum += max(ccc_list) -# return max_ccc_sum - - -# def compute_FRC(args, pred, listener_em, val_test='val'): -# pred = pred -# listener_em = listener_em -# if val_test == 'val': -# speaker_neighbour_matrix = np.load(os.path.join(args.dataset_path, 'neighbour_emotion_val.npy')) -# # speaker_neighbour_matrix = np.load(os.path.join(args.dataset_path, 'person_specific_masked_neighbour_emotion_val.npy')) -# else: -# speaker_neighbour_matrix = np.load(os.path.join(args.dataset_path, 'neighbour_emotion_test.npy')) -# # speaker_neighbour_matrix = np.load(os.path.join(args.dataset_path, 'person_specific_masked_neighbour_emotion_test.npy')) -# -# all_FRC_list = [] -# for i in range(pred.shape[1]): -# FRC_list = [] -# for k in range(pred.shape[0]): -# speaker_neighbour_index = np.argwhere(speaker_neighbour_matrix[k] == 1).reshape(-1) -# speaker_neighbour_index_len = len(speaker_neighbour_index) -# ccc_list = [] -# for n_index in range(speaker_neighbour_index_len): -# ''' -# listener_em order :[listener1, listener2, listener3, ....., listener_n, speaker1, speaker2, speaker3, ....., speaker_n] -# listener1: [1, emotion_dim] -# -# listener_em[speaker_neighbour_index[n_index]]: -# 1. speaker_neighbour_index[n_index]: speaker_j (with similar emotion as the speaker_k) -# 2. listener_em[speaker_neighbour_index[n_index]]: emotion features of listener_j (speaker_j -> listener_j) -# So we can get an additional GT listener_j to listener_k (i.e., speaker_j -> listener_k) -# ''' -# -# similar_listener_emotion = listener_em[speaker_neighbour_index[n_index]] -# ccc = concordance_correlation_coefficient(similar_listener_emotion.numpy(), pred[k, i].numpy()) -# ccc_list.append(ccc) -# max_ccc = max(ccc_list) -# FRC_list.append(max_ccc) -# all_FRC_list.append(np.mean(FRC_list)) -# return sum(all_FRC_list) diff --git a/personalised/code/framework/metrics/FRD.py b/personalised/code/framework/metrics/FRD.py deleted file mode 100644 index c5ee673fa2f1d5bcd4b6ad116677fae720193529..0000000000000000000000000000000000000000 --- a/personalised/code/framework/metrics/FRD.py +++ /dev/null @@ -1,45 +0,0 @@ -import numpy as np -from tslearn.metrics import dtw -from functools import partial -import multiprocessing as mp -from tqdm import tqdm - - -def _func(target, pred): - target = target.numpy() - pred = pred.numpy() - - # num_preds = pred.shape[0] - min_dwt_sum = 0 - for i in range(pred.shape[0]): - dwt_list = [] - for j in range(target.shape[0]): - emotion = target[j] - res = 0 - for st, ed, weight in [(0, 15, 1 / 15), (15, 17, 1), (17, 25, 1 / 8)]: - res += weight * dtw(pred[i].astype(np.float32)[:, st: ed], emotion.astype(np.float32)[:, st: ed]) - dwt_list.append(res) - min_dwt_sum += min(dwt_list) - return min_dwt_sum - - -def _func_star(args): - target, pred = args - return _func(target, pred) - - -def compute_FRD(preds, targets, p=1): - tasks = list(zip(targets, preds)) - FRD_list = [] - - with mp.Pool(processes=p) as pool: - for result in tqdm( - pool.imap_unordered(_func_star, tasks), - total=len(tasks), - desc="Computing FRD" - ): - FRD_list.append(result) - - return np.mean(FRD_list) - - diff --git a/personalised/code/framework/metrics/FRDvs.py b/personalised/code/framework/metrics/FRDvs.py deleted file mode 100644 index 5285bba82ba24ae38734602019a3372cfa6dd26f..0000000000000000000000000000000000000000 --- a/personalised/code/framework/metrics/FRDvs.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch -from einops import rearrange - - -def compute_FRDvs(preds): - NotImplemented - - -# def compute_FRDvs(preds): -# # preds: List: [Tensor([num_preds, l, 25]), Tensor([num_preds, l', 25]), ] -# group_scores = [] -# for pred_item in preds: -# # num_preds, L, C = pred_item.shape -# flat = rearrange(pred_item, 'n l c -> n (l c)') -# dist = torch.pow(torch.cdist(flat, flat), 2) # (num_preds, num_preds) -# preds_ = preds.reshape(preds.shape[0], preds.shape[1], -1) -# preds_ = preds_.transpose(0, 1) -# # preds_: (10, N, 750*...) -# dist = torch.pow(torch.cdist(preds_, preds_), 2) -# # dist: (10, N, N) -# dist = torch.sum(dist) / (preds.shape[0] * (preds.shape[0] - 1) * preds.shape[1]) -# return dist / preds_.shape[-1] - - -# def compute_FRDvs(preds): -# # preds: (N, 10, 750, ...) -# preds_ = preds.reshape(preds.shape[0], preds.shape[1], -1) -# preds_ = preds_.transpose(0, 1) -# # preds_: (10, N, 750*...) -# dist = torch.pow(torch.cdist(preds_, preds_), 2) -# # dist: (10, N, N) -# dist = torch.sum(dist) / (preds.shape[0] * (preds.shape[0] - 1) * preds.shape[1]) -# return dist / preds_.shape[-1] \ No newline at end of file diff --git a/personalised/code/framework/metrics/FRVar.py b/personalised/code/framework/metrics/FRVar.py deleted file mode 100644 index 2333548fb622937985a7010b7c8b4e1739072739..0000000000000000000000000000000000000000 --- a/personalised/code/framework/metrics/FRVar.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch -from tqdm import tqdm - - -def compute_FRVar(preds): - # preds: [Tensor(num_preds, l, 25), Tensor(num_preds, l, 25), ...] - variance = [] - for pred in tqdm(preds, desc="Computing FRVar"): - - variance.append(torch.var(pred, dim=1)) - - variance = torch.stack(variance, dim=0) - return torch.mean(variance) - - -# def compute_FRVar(preds): -# if len(preds.shape) == 3: -# # preds: (10, 750, ...) -# var = torch.var(preds, dim=1) -# return torch.mean(var) -# elif len(preds.shape) == 4: -# # preds: (N, 10, 750, ...) -# var = torch.var(preds, dim=2) -# return torch.mean(var) diff --git a/personalised/code/framework/metrics/S_MSE.py b/personalised/code/framework/metrics/S_MSE.py deleted file mode 100644 index 5cdbbdfaaf731d0fd775fc958df2431a79b63e18..0000000000000000000000000000000000000000 --- a/personalised/code/framework/metrics/S_MSE.py +++ /dev/null @@ -1,25 +0,0 @@ -import torch -from tqdm import tqdm - - -def compute_s_mse(preds): - # preds: List: [Tensor([num_preds, l, 25]), Tensor([num_preds, l', 25]), ...] - - dist = 0 - for pred_item in tqdm(preds, desc="Computing S_MSE"): - assert pred_item.shape[0] > 1, "num_preds set to greater than 1" - pred_item_ = pred_item.reshape(pred_item.shape[0], -1) # (num_preds, d) - dist_ = torch.pow(torch.cdist(pred_item_, pred_item_), 2) # (num_preds, num_preds) - dist_ = torch.sum(dist_) / (pred_item_.shape[0] * (pred_item_.shape[0] - 1) * pred_item_.shape[1]) - dist += dist_ - return dist / len(preds) - - # for b in range(preds.shape[0]): - # preds_item = preds[b] - # if preds_item.shape[0] == 1: - # return 0.0 - # preds_item_ = preds_item.reshape(preds_item.shape[0], -1) - # dist_ = torch.pow(torch.cdist(preds_item_, preds_item_), 2) - # dist_ = torch.sum(dist_) / (preds_item.shape[0] * (preds_item.shape[0] - 1) * preds_item_.shape[1]) - # dist += dist_ - # return dist / preds.shape[0] \ No newline at end of file diff --git a/personalised/code/framework/metrics/TLCC.py b/personalised/code/framework/metrics/TLCC.py deleted file mode 100644 index 377868a309c7f279df8cf0bbdaa58cd22d9fcbff..0000000000000000000000000000000000000000 --- a/personalised/code/framework/metrics/TLCC.py +++ /dev/null @@ -1,101 +0,0 @@ -import torch -import numpy as np -import multiprocessing as mp -from tqdm import tqdm - - -def crosscorr(datax, datay, lag=0, dim=25): - pcc_list = [] - for i in range(dim): - cn_1, cn_2 = shift(datax[:, i], datay[:, i], lag) - pcc_i = np.corrcoef(cn_1, cn_2)[0, 1] - # pcc_i = torch.corrcoef(torch.stack([cn_1, cn_2], dim=0).float())[0, 1] - pcc_list.append(pcc_i.item()) - return torch.mean(torch.Tensor(pcc_list)) - - -def calculate_tlcc(pred, sp, seconds=2, fps=25): - rs = [crosscorr(pred, sp, lag, sp.shape[-1]) for lag in range(-int(seconds * fps - 1), int(seconds * fps))] - peak = max(rs) - center = rs[len(rs) // 2] - offset = len(rs) // 2 - torch.argmax(torch.Tensor(rs)) - return peak, center, offset - - -def _func(pred_item, sp_item): - # pred_item: Tensor([num_preds, l, 25]) - # sp_item: Tensor([l, 25]) - min_len = min(pred_item.shape[1], sp_item.shape[0]) - pred_item = pred_item[:, :min_len, :] - sp_item = sp_item[:min_len, :] - for i in range(pred_item.shape[0]): - peak, center, offset = calculate_tlcc( - pred_item[i].float(), - sp_item.float(), - ) - return torch.abs(offset).item() - - -def _func_star(args): - pred_item, sp_item = args - return _func(pred_item, sp_item) - - -def compute_TLCC(preds, speakers, p=8): - # preds: List: [Tensor([num_preds, l, 25]), Tensor([num_preds, l', 25]), ...] - # speakers: List: [Tensor([l, 25]), Tensor([l', 25]), ...] - tasks = list(zip(preds, speakers)) - offset_list = [] - - np.seterr(divide='ignore', invalid='ignore') - - with mp.Pool(processes=p) as pool: - for result in tqdm( - pool.imap(_func_star, tasks), - total=len(tasks), - desc="Computing TLCC" - ): - offset_list.append(result) - - return torch.mean(torch.Tensor(offset_list)).item() - - -def shift(x, y, lag): - if lag > 0: - return x[lag:], y[:-lag] - elif lag < 0: - return x[:lag], y[-lag:] - else: - return x, y - - -# def _func(pred_item, sp_item): -# for i in range(pred_item.shape[0]): -# peak, center, offset = calculate_tlcc(pred_item[i], sp_item) -# return torch.abs(offset).item() - - -# def compute_TLCC(pred, speaker): -# # pred: N 10 750 25 -# # speaker: N 750 25 -# offset_list = [] -# for k in range(speaker.shape[0]): -# pred_item = pred[k] -# sp_item = speaker[k] -# for i in range(pred_item.shape[0]): -# peak, center, offset = calculate_tlcc(pred_item[i].float().numpy(), sp_item.float().numpy()) -# offset_list.append(torch.abs(offset).item()) -# return torch.mean(torch.Tensor(offset_list)).item() - - -# def compute_TLCC_mp(pred, speaker, p=8): -# # pred: N 10 750 dim -# # speaker: N 750 dim -# offset_list = [] -# # process each speaker in parallel -# np.seterr(divide='ignore', invalid='ignore') -# -# with mp.Pool(processes=p) as pool: -# # use map -# offset_list += pool.starmap(_func, zip(pred.float().numpy(), speaker.float().numpy())) -# return torch.mean(torch.Tensor(offset_list)).item() diff --git a/personalised/code/framework/metrics/__init__.py b/personalised/code/framework/metrics/__init__.py deleted file mode 100644 index 94bb02c15266d4dc5f328140d8ccdf037426c991..0000000000000000000000000000000000000000 --- a/personalised/code/framework/metrics/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from .FRC import compute_FRC -from .FRD import compute_FRD -from .FRDvs import compute_FRDvs -from .FRVar import compute_FRVar -from .S_MSE import compute_s_mse -from .TLCC import compute_TLCC diff --git a/personalised/code/framework/metrics/metric.py b/personalised/code/framework/metrics/metric.py deleted file mode 100644 index 9cb8995a5139a7941f0065ff4d71c8062123280c..0000000000000000000000000000000000000000 --- a/personalised/code/framework/metrics/metric.py +++ /dev/null @@ -1,90 +0,0 @@ -import torch -import numpy as np -from einops import repeat - - -def sample_truncnorm(shape, mean=0.0, std=0.5, low=-1.0, high=1.0): - samples = torch.empty(shape, dtype=torch.float32) - mask = torch.ones(shape, dtype=torch.bool) - while mask.any(): - new = torch.randn(shape, dtype=torch.float32) * std + mean - samples[mask] = new[mask] - mask = (samples < low) | (samples > high) - - return samples - - -def baseline_random(target): - B, L, D = target.shape - assert D == 25 - - bin_part = torch.randint(0, 2, (B, L, 15), dtype=torch.float32) - # cont_part = torch.empty((B, L, 2)).uniform_(-1.0, 1.0) - cont_part = sample_truncnorm((B, L, 2), mean=0.0, std=0.5, low=-1.0, high=1.0) - cat_raw = torch.rand((B, L, 8)) - cat_part = cat_raw / cat_raw.sum(dim=-1, keepdim=True) - - return torch.cat([bin_part, cont_part, cat_part], dim=-1) - - -def baseline_mime(target): - return repeat(target, 'l d -> b l d', b=10) - - -# def baseline_meanseq(target, pred=None): -# # mean over time dim=1 → shape (B,1,25) -# mean_seq = target.mean(dim=1, keepdim=True) -# return mean_seq.repeat(1, target.size(1), 1) - - -# def baseline_meanfr(target, pred=None): -# # mean over batch and time dims → shape (25,) -# global_mean = target.mean(dim=(0,1), keepdim=True) # → (1,1,25) -# B, L, _ = target.shape -# return global_mean.expand(B, L, target.size(2)) - - -def s_mse(preds): - # preds: (B, 10, 750, 25) - dist = 0 - for b in range(preds.shape[0]): - preds_item = preds[b] - if preds_item.shape[0] == 1: - return 0.0 - preds_item_ = preds_item.reshape(preds_item.shape[0], -1) - dist_ = torch.pow(torch.cdist(preds_item_, preds_item_), 2) - dist_ = torch.sum(dist_) / (preds_item.shape[0] * (preds_item.shape[0] - 1) * preds_item_.shape[1]) - dist += dist_ - return dist / preds.shape[0] - - -def FRVar(preds): - if len(preds.shape) == 3: - # preds: (10, 750, 25) - var = torch.var(preds, dim=1) - return torch.mean(var) - elif len(preds.shape) == 4: - # preds: (N, 10, 750, 25) - var = torch.var(preds, dim=2) - return torch.mean(var) - - -def FRDvs(preds): - # preds: (N, 10, 750, 25) - preds_ = preds.reshape(preds.shape[0], preds.shape[1], -1) - preds_ = preds_.transpose(0, 1) - # preds_: (10, N, 750*25) - dist = torch.pow(torch.cdist(preds_, preds_), 2) - # dist: (10, N, N) - dist = torch.sum(dist) / (preds.shape[0] * (preds.shape[0] - 1) * preds.shape[1]) - return dist / preds_.shape[-1] - - -def compute_FRVar(pred): - FRVar_list = [] - for k in range(pred.shape[0]): - pred_item = pred[k] - for i in range(0, pred_item.shape[0]): - var = np.mean(np.var(pred_item[i].numpy().astype(np.float32), axis=0)) - FRVar_list.append(var) - return np.mean(FRVar_list) diff --git a/personalised/code/framework/modules/__init__.py b/personalised/code/framework/modules/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/personalised/code/framework/modules/emotion_autoencoder.py b/personalised/code/framework/modules/emotion_autoencoder.py deleted file mode 100644 index b4e0fe1c062762924106999e245a2e380cdf2277..0000000000000000000000000000000000000000 --- a/personalised/code/framework/modules/emotion_autoencoder.py +++ /dev/null @@ -1,310 +0,0 @@ -import torch -import torch.nn as nn -from einops import repeat -import numpy as np -from framework.motion_diffusion.diffusion.utils.temos_utils import lengths_to_mask - - -def sequence_slice(emb, start_indices, end_indices, max_seq_length): - """ - :param emb: positional token embeddings | time query token embeddings - (bz, l=5000, d_model) - """ - B, L, D = emb.shape - - lengths = end_indices - start_indices + 1 # [B] - M = int(lengths.max()) # Python int - - rel_pos = torch.arange(M, device=emb.device).unsqueeze(0).expand(B, M) # [B, M] - abs_pos = start_indices.unsqueeze(1) + rel_pos # [B, M] - abs_pos = abs_pos.clamp(0, L - 1) # will mask out if out of bounds - - slice_batched = emb.gather( - dim=1, - index=abs_pos.unsqueeze(-1).expand(-1, -1, D) - ) - - mask = rel_pos < lengths.unsqueeze(1) - slice_batched = slice_batched * mask.unsqueeze(-1) # [B, M, D] - - assert M <= max_seq_length, "Length of sliced sequence exceeds max_seq_length" - slice_batched = torch.cat((slice_batched, - torch.zeros(B, max_seq_length - M, D, device=emb.device)), dim=1) - - return slice_batched # lengths - - -class PositionalEncoding(nn.Module): - def __init__(self, d_model, dropout = 0.1, pe_type = "absolute", - batch_first = True, max_len = 50000,): - super().__init__() - self.batch_first = batch_first - self.dropout = nn.Dropout(p=dropout) - self.pe_type = pe_type - - if pe_type == "learnable": - self.pe = nn.Parameter(torch.zeros(1, max_len, d_model)) - nn.init.uniform_(self.pe) - elif pe_type == "absolute": - pe = torch.zeros(max_len, d_model) - position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) - div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model)) - pe[:, 0::2] = torch.sin(position * div_term) - pe[:, 1::2] = torch.cos(position * div_term) - pe = pe.unsqueeze(0).transpose(0, 1) - self.register_buffer('pe', pe) - else: - raise ValueError("Unknown positional encoding type: {}".format(pe_type)) - - def forward(self, x, start_indices, end_indices, clip_length): - if self.pe_type == "learnable": - pe = repeat(self.pe, "1 l d -> b l d", b=x.shape[0]) - elif self.pe_type == "absolute": - assert self.batch_first, "At the moment only batch_first=True" - pe = repeat(self.pe.permute(1, 0, 2), "1 l d -> b l d", b=x.shape[0]) - else: - raise ValueError("Unknown positional encoding type: {}".format(self.pe_type)) - - pe = sequence_slice(pe, start_indices, end_indices, clip_length) - x = x + pe[:, :x.shape[1], :] - return self.dropout(x) - - -class Encoder(nn.Module): - def __init__(self, - d_model = 512, - nhead = 8, - num_layers = 6, - max_seq_len = 5000, - global_token_len = 128, - mlp_dist = False, - ): - super().__init__() - self.max_seq_len = max_seq_len - - if mlp_dist: - self.global_token_len = global_token_len - else: - self.global_token_len = global_token_len * 2 - self.global_tokens = nn.Parameter(torch.zeros(1, self.global_token_len, d_model)) - self.reset_parameters() - - self.transformer_encoder = nn.TransformerEncoder( - nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True), - num_layers=num_layers) - - def reset_parameters(self): - nn.init.normal_(self.global_tokens, mean=0.0, std=0.04) # 0.02? - - def forward(self, - x: torch.Tensor, - pe: nn.Module, - start_indices: torch.Tensor, - end_indices: torch.Tensor, - clip_length: torch.Tensor, - ): - - lengths = end_indices - start_indices + 1 - x = pe(x, start_indices, end_indices, clip_length) - x = torch.cat((self.global_tokens.expand(x.shape[0], -1, -1), x), dim=1) - mask = lengths_to_mask(lengths, device=x.device, max_len=clip_length) - mask = torch.cat((torch.ones(x.shape[0], self.global_token_len, device=x.device), mask), dim=-1) - src_key_padding_mask = ~(mask.bool()) - - x = self.transformer_encoder(src=x, mask=None, src_key_padding_mask=src_key_padding_mask) - return x[:, :self.global_token_len, :] - - -class Decoder(nn.Module): - def __init__(self, - d_model = 512, - nhead = 8, - num_layers = 6, - max_seq_len = 5000, - query_type = "learnable"): - super().__init__() - self.d_model = d_model - self.max_seq_len = max_seq_len - self.query_type = query_type - - if query_type == "learnable": - self.time_query = nn.Parameter(torch.zeros(1, self.max_seq_len, d_model)) - # torch.nn.init.normal_(self.time_query, mean=0.0, std=0.04) - nn.init.uniform_(self.time_query) - - self.transformer_decoder = nn.TransformerDecoder( - nn.TransformerDecoderLayer(d_model=d_model, nhead=nhead, batch_first=True), - num_layers=num_layers) - - def forward(self, x: torch.Tensor, - pe: nn.Module, - start_indices: torch.Tensor, - end_indices: torch.Tensor, - clip_length: torch.Tensor, - ): - - lengths = end_indices - start_indices + 1 - mask = lengths_to_mask(lengths, device=x.device, max_len=clip_length) - - if self.query_type == "learnable": - time_query = sequence_slice(repeat(self.time_query, "1 l d -> b l d", b=x.shape[0]), - start_indices, end_indices, clip_length) - elif self.query_type == "zero": - time_query = torch.zeros(x.shape[0], clip_length, self.d_model).to(x.device) - else: - raise ValueError("Unknown query type: {}".format(self.query_type)) - # (bsz, clip_len, d) - time_query = pe(time_query, start_indices, end_indices, clip_length) - - x = self.transformer_decoder(tgt=time_query, memory=x, tgt_key_padding_mask=~mask) - # padding_mask = mask.float().unsqueeze(-1).to(x.device) - # x = x * padding_mask - - return x - - -class EmotionVAE(nn.Module): - def __init__(self, - in_channels = 25, - out_channels = 25, - feature_dim = 512, - nhead = 8, - dropout = 0.1, - num_encoder_layers = 6, - num_decoder_layers = 6, - mlp_dist = False, # expand mu & logvar - in_proj_type = "linear", # linear | mlp - out_proj_type = "separate", # separate | shared - pe_type = "learnable", - query_type = "zero", - max_seq_len = 5000, - global_token_len = 128, - **kwargs, - ): - - super().__init__() - self.feature_dim = feature_dim - self.mlp_dist = mlp_dist - self.in_proj_type = in_proj_type - self.out_proj_type = out_proj_type - self.pe_type = pe_type - self.global_token_len = global_token_len - - self.PE = PositionalEncoding(d_model=feature_dim, - dropout=dropout, - pe_type=pe_type, - max_len=max_seq_len) - - if self.in_proj_type == "mlp": - self.in_proj = nn.Sequential( - nn.Linear(in_channels, feature_dim), - nn.LayerNorm(feature_dim), - nn.GELU(), - nn.Linear(feature_dim, feature_dim), - nn.LayerNorm(feature_dim), - ) - elif self.in_proj_type == "linear": - self.in_proj = nn.Linear(in_channels, feature_dim) - else: - raise ValueError("Unknown input projection type: {}".format(self.in_proj_type)) - - if self.out_proj_type == "separate": - self.au_out_proj = nn.Linear(feature_dim, 15) - self.va_out_proj = nn.Linear(feature_dim, 2) - self.emo_out_proj = nn.Linear(feature_dim, 8) - elif self.out_proj_type == "shared": - self.out_proj = nn.Linear(feature_dim, out_channels) - else: - raise ValueError("Unknown output projection type: {}".format(self.out_proj_type)) - - self.encoder = Encoder( - d_model=feature_dim, - nhead=nhead, - num_layers=num_encoder_layers, - max_seq_len=max_seq_len, - global_token_len=global_token_len, - mlp_dist=mlp_dist, - ) - - if mlp_dist: - self.mu_head = nn.Linear(feature_dim, feature_dim) - self.logvar_head = nn.Linear(feature_dim, feature_dim) - - self.decoder = Decoder( - d_model=feature_dim, - num_layers=num_decoder_layers, - max_seq_len=max_seq_len, - query_type=query_type, - ) - - def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor, deterministic: bool = False): - """ - Reparameterization trick to sample from N(mu, var) from - N(0,1). - :param mu: (Tensor) Mean of the latent Gaussian [B x M x D] - :param logvar: (Tensor) Standard deviation of the latent Gaussian [B x M x D] - """ - # std = torch.exp(0.5 * logvar) - # eps = torch.randn_like(std) - - if deterministic: - return mu, None - - std = logvar.exp().pow(0.5) - dist = torch.distributions.Normal(mu, std) - latent = self.sample_from_distribution(dist).to(mu.device) - return latent, dist - - def sample_from_distribution(self, distribution): - return distribution.rsample() - - def forward(self, x: torch.Tensor, - start_e: torch.Tensor, - end_e: torch.Tensor, - start_d: torch.Tensor, - end_d: torch.Tensor, - reparameterization: str = "random"): - """ - :param x: input sequence, (batch_size, token_len, feature_dim) - """ - B, L, D = x.shape - x = self.in_proj(x) - z = self.encoder(x, self.PE, start_e, end_e, L) # L: seq_len - - if self.mlp_dist: - mu = self.mu_head(z) - logvar = self.logvar_head(z) - else: - mu = z[:, :self.global_token_len, :] - logvar = z[:, self.global_token_len:, :] - - deterministic = reparameterization == "deterministic" - latent, dist = self.reparameterize(mu, logvar, deterministic=deterministic) - - x = self.decoder(latent, self.PE, start_d, end_d, L) - - if self.out_proj_type == "shared": - x = self.out_proj(x) - au_out = x[:, :, :15] # F.sigmoid(x[:, :, :15]) - va_out = x[:, :, 15:17] - emo_out = x[:, :, 17:] # F.softmax(x[:, :, 17:], dim=-1) - out = (au_out, va_out, emo_out) # torch.cat((au_out, va_out, emo_out), dim=-1) - elif self.out_proj_type == "separate": - au_out = self.au_out_proj(x) # F.sigmoid(self.au_out_proj(x)) - va_out = self.va_out_proj(x) - emo_out = self.emo_out_proj(x) # F.softmax(self.emo_out_proj(x), dim=-1) - out = (au_out, va_out, emo_out) # torch.cat((au_out, va_out, emo_out), dim=-1) - else: - raise ValueError("Unknown output projection type: {}".format(self.out_proj_type)) - - out_padding_mask = lengths_to_mask( - lengths=(end_d - start_d + 1), - device=x.device, - max_len=L, - ).float().unsqueeze(-1) - out = (o * out_padding_mask for o in out) - - return out, latent, dist, out_padding_mask - - def get_model_name(self): - return self.__class__.__name__ \ No newline at end of file diff --git a/personalised/code/framework/modules/post_processor.py b/personalised/code/framework/modules/post_processor.py deleted file mode 100644 index c00c62a4dc6b6cc01f1d1fbe4715fc5624d4cbf5..0000000000000000000000000000000000000000 --- a/personalised/code/framework/modules/post_processor.py +++ /dev/null @@ -1,149 +0,0 @@ -from pathlib import Path -from typing import Optional -import torch -import torch.nn.functional as F -import math -from einops import rearrange, repeat -import hydra -from omegaconf import OmegaConf -from hydra.utils import instantiate -import os -from framework.utils.util import from_pretrained_checkpoint - - -class Processor: - def __init__(self, - config_name: str = "configs/shared/model/emotion_autoencoder.yaml", - ckpt_dir: str = "pretrained_models/post_processor", - device: Optional[torch.device] = None, - clip_len_test: int = 1000, - cfg_dir: str = None, - **kwargs): - if cfg_dir is None: - cfg_dir = hydra.utils.get_original_cwd() - cfg = OmegaConf.load(os.path.join(cfg_dir, config_name)) - self.model = instantiate(cfg, _recursive_=False) - self.clip_len_test = clip_len_test - self.num_preds = kwargs.get("num_preds", 10) - - if device is None: - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - self.device = device - - ckpt_path = self.get_ckpt_path(ckpt_dir) - from_pretrained_checkpoint(ckpt_path, self.model, device=device) - self.model.eval() - - def get_ckpt_path(self, ckpt_dir): - ckpt_path = Path(hydra.utils.to_absolute_path(ckpt_dir)) / 'checkpoint.pth' - assert ckpt_path.is_file(), f"Checkpoint file not found at {ckpt_dir}" - return ckpt_path - - def forward(self, prediction_list, target_list): - processed_target_list = [] - - for predictions, targets in zip(prediction_list, target_list): - if len(predictions.shape) == 2: - predictions = repeat(predictions, 'l d -> n l d', n=self.num_preds) - - # predictions: Tensor([num_preds, l, d]) - # targets: List: [(l', d), (l'', d), ...] - _, pred_seq_len, dim = predictions.shape - - processed_targets = [] - for tgt in targets: # tgt: Tensor([l', 25]) - tgt_seq_len = tgt.shape[0] - max_len = self.clip_len_test - - if pred_seq_len == tgt_seq_len: - processed_targets.append(tgt) - continue - - if pred_seq_len > tgt_seq_len: - num_segments = math.ceil(pred_seq_len / max_len) - assert tgt_seq_len >= num_segments - min_len = math.ceil(tgt_seq_len / num_segments) - total_len = num_segments * min_len - - tgt = torch.cat( - (tgt, - torch.zeros(size=(int(total_len - tgt_seq_len), dim)) - ), dim=0) - tgt = rearrange(tgt, '(b l) d -> b l d', b=num_segments) - tgt = torch.cat((tgt, torch.zeros(num_segments, max_len - min_len, dim)), dim=1) - - out_start_indices = torch.zeros(size=(num_segments,)) - input_start_indices = torch.zeros(size=(num_segments,)) - if num_segments == 1: - out_end_indices = torch.tensor([pred_seq_len - 1]) - input_end_indices = torch.tensor([tgt_seq_len - 1]) - else: - if pred_seq_len % max_len == 0: - out_end_indices = torch.tensor([max_len - 1] * num_segments) - else: - out_end_indices = torch.cat((torch.tensor([max_len - 1] * (num_segments - 1)), - torch.tensor([pred_seq_len % max_len - 1]))) - if tgt_seq_len % min_len == 0: - input_end_indices = torch.tensor([min_len - 1] * num_segments) - else: - input_end_indices = torch.cat((torch.tensor([min_len - 1] * (num_segments - 1)), - torch.tensor([tgt_seq_len % min_len - 1]))) - - else: - num_segments = math.ceil(tgt_seq_len / max_len) - assert pred_seq_len >= num_segments - total_len = num_segments * max_len - min_len = math.ceil(pred_seq_len / num_segments) - - tgt = torch.cat( - (tgt, - torch.zeros(size=(int(total_len - tgt_seq_len), dim)) - ), dim=0) - tgt = rearrange(tgt, '(b l) d -> b l d', b=num_segments) - - out_start_indices = torch.zeros(size=(num_segments,)) - input_start_indices = torch.zeros(size=(num_segments,)) - if num_segments == 1: - out_end_indices = torch.tensor([pred_seq_len - 1]) - input_end_indices = torch.tensor([tgt_seq_len - 1]) - else: - if tgt_seq_len % max_len == 0: - input_end_indices = torch.tensor([max_len - 1] * num_segments) - else: - input_end_indices = torch.cat((torch.tensor([max_len - 1] * (num_segments - 1)), - torch.tensor([tgt_seq_len % max_len - 1]))) - if pred_seq_len % min_len == 0: - out_end_indices = torch.tensor([min_len - 1] * num_segments) - else: - out_end_indices = torch.cat((torch.tensor([min_len - 1] * (num_segments - 1)), - torch.tensor([pred_seq_len % min_len - 1]))) - - lengths = (out_end_indices - out_start_indices + 1).long() - inputs, input_start_indices, input_end_indices, out_start_indices, out_end_indices = \ - tgt.to(self.device), input_start_indices.to(self.device), input_end_indices.to(self.device), \ - out_start_indices.to(self.device), out_end_indices.to(self.device) - - outputs = self.model( - inputs, - input_start_indices.long(), - input_end_indices.long(), - out_start_indices.long(), - out_end_indices.long(), - )[0] # (bsz, total_len, d) - - processed_target = [] - for i, (out_au, out_va, out_em) in enumerate(zip(*outputs)): - _len = lengths[i] - out_au = (F.sigmoid(out_au) >= 0.5).float() - out_em = F.softmax(out_em, dim=-1).float() - out_all = torch.cat((out_au, out_va, out_em), dim=-1)[:_len].detach().cpu() - processed_target.append(out_all) - processed_target = torch.cat(processed_target, dim=0) - processed_targets.append(processed_target) # len equal to speaker's - - processed_targets = torch.stack(processed_targets, dim=0) - processed_target_list.append(processed_targets) - # prediction_list, processed_target_list - # List [Tensor([num_preds, l, 25]), Tensor([num_preds, l', 25]), ...] - - return processed_target_list diff --git a/personalised/code/framework/motion_diffusion/diffusion/__init__.py b/personalised/code/framework/motion_diffusion/diffusion/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/__init__.py b/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/transformer_denoiser.py b/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/transformer_denoiser.py deleted file mode 100644 index 33f25255110cccb9927947118c803530ab64db2b..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/transformer_denoiser.py +++ /dev/null @@ -1,562 +0,0 @@ -from typing import Dict, List -import math -import torch -import torch.nn as nn -from torch import Tensor -# from framework.motion_diffusion.diffusion.utils.util import tgt_biased_mask, memory_biased_mask -from framework.motion_diffusion.diffusion.operator.embeddings import (TimestepEmbedding, Timesteps) -from framework.motion_diffusion.diffusion.operator.position_encoding import build_position_encoding -from framework.motion_diffusion.diffusion.operator.cross_attention import (SkipTransformerEncoder, - TransformerDecoder, - TransformerDecoderLayer, - TransformerEncoderLayer) - - -def lengths_to_mask(lengths: List[int], - device: torch.device, - max_len: int = None) -> Tensor: - lengths = torch.tensor(lengths, device=device) - max_len = max_len if max_len else max(lengths) - mask = torch.arange(max_len, device=device).expand( - len(lengths), max_len) < lengths.unsqueeze(1) - return mask - - -def timestep_embedding(timesteps, dim, max_period=10000, dtype=torch.float32): - """ - Create sinusoidal timestep embeddings. - :param timesteps: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param dim: the dimension of the output. - :param max_period: controls the minimum frequency of the embeddings. - :return: an [N x dim] Tensor of positional embeddings. - """ - half = dim // 2 - freqs = torch.exp( - -math.log(max_period) * torch.arange(start=0, end=half, dtype=dtype) / half - ).to(device=timesteps.device) - args = timesteps[:, None].type(dtype) * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) - return embedding - - -class TransformerDenoiser(nn.Module): - def __init__(self, - task: str = 'offline', - window_size: int = 30, - encode_emotion: bool = False, - encode_3dmm: bool = False, - ablation_skip_connection: bool = True, - nfeats: int = 25, - latent_dim: int = 512, - ff_size: int = 1024, - num_layers: int = 7, - num_heads: int = 4, - dropout: float = 0.1, - normalize_before: bool = False, - activation: str = "gelu", - flip_sin_to_cos: bool = True, - return_intermediate_dec: bool = False, - position_embedding: str = "learned", - arch: str = "trans_dec", - freq_shift: int = 0, - time_encoded_dim: int = 64, - s_audio_dim: int = 78, - s_audio_scale: float = 1.0, - s_emotion_dim: int = 25, - s_embed_dim: int = 512, - s_3dmm_dim: int = 58, - concat: str = "concat_first", - condition_concat: str = "token_concat", - guidance_scale: float = 7.5, - s_audio_enc_drop_prob: float = 0.2, # speaker_audio_encodings - s_latent_embed_drop_prob: float = 0.2, # speaker_latent_embed - s_3dmm_enc_drop_prob: float = 0.2, # speaker_3dmm_encodings - s_emotion_enc_drop_prob: float = 1.0, # speaker_emotion_encodings - past_l_emotion_drop_prob: float = 1.0, # past_listener_emotion - **kwargs) -> None: - super().__init__() - - self.task = task - self.window_size = window_size - self.encode_emotion = encode_emotion - self.encode_3dmm = encode_3dmm - self.s_audio_scale = s_audio_scale - self.latent_dim = latent_dim - self.ablation_skip_connection = ablation_skip_connection - self.arch = arch - self.concat = concat - self.condition_concat = condition_concat - # for classifier-free guidance - self.guidance_scale = guidance_scale - self.s_audio_enc_drop_prob = s_audio_enc_drop_prob - self.s_latent_embed_drop_prob = s_latent_embed_drop_prob - self.s_3dmm_enc_drop_prob = s_3dmm_enc_drop_prob - self.s_emotion_enc_drop_prob = s_emotion_enc_drop_prob - self.past_l_emotion_drop_prob = past_l_emotion_drop_prob - - # project between emotion output feat and emotion latent embedding - self.to_emotion_embed = nn.Linear(nfeats, self.latent_dim) if nfeats != self.latent_dim else nn.Identity() - self.to_emotion_feat = nn.Linear(self.latent_dim, nfeats) if self.latent_dim != nfeats else nn.Identity() - - # project time to latent_dim - self.time_proj = Timesteps(time_encoded_dim, flip_sin_to_cos, freq_shift) - self.time_embedding = TimestepEmbedding(time_encoded_dim, self.latent_dim) - - self.speaker_latent_proj = nn.Sequential(nn.ReLU(), nn.Linear(s_embed_dim, self.latent_dim)) \ - if s_embed_dim != self.latent_dim else nn.Identity() - - self.speaker_audio_proj = nn.Linear(s_audio_dim, self.latent_dim) \ - if s_audio_dim != self.latent_dim else nn.Identity() - - if self.encode_3dmm: # assume dimension of encoded 3dmm equals latent_dim - self.speaker_3dmm_proj = nn.Identity() - else: - assert s_3dmm_dim != self.latent_dim, "wrong dimension of raw 3dmm features." - self.speaker_3dmm_proj = nn.Linear(s_3dmm_dim, self.latent_dim) - - if self.encode_emotion: # assume dimension of encoded emotion equals latent_dim - self.emotion_proj = nn.Identity() - else: - assert s_emotion_dim != self.latent_dim, "wrong dimension of raw emotion features." - self.emotion_proj = nn.Linear(s_emotion_dim, self.latent_dim) - - self.query_pos = build_position_encoding( - self.latent_dim, position_embedding=position_embedding) - self.mem_pos = build_position_encoding( - self.latent_dim, position_embedding=position_embedding) - - # we concat conditions (including: speaker 3dmm, speaker audio, speaker emotion encodings) along last dimension. - self.condition_proj = nn.Linear(self.latent_dim * 3, self.latent_dim) \ - if self.condition_concat == 'feat_concat' else nn.Identity() - - # define our transformer decoder layer - decoder_layer = TransformerDecoderLayer( - self.latent_dim, - num_heads, - ff_size, - dropout, - activation, - normalize_before, - ) - - if self.arch == "trans_enc": # Transformer Encoder - if self.ablation_skip_connection: - # use DETR transformer - encoder_layer = TransformerEncoderLayer( - self.latent_dim, - num_heads, - ff_size, - dropout, - activation, - normalize_before, - ) - encoder_norm = nn.LayerNorm(self.latent_dim) - self.encoder = SkipTransformerEncoder(encoder_layer, - num_layers, encoder_norm) - else: - # use torch transformer - encoder_layer = nn.TransformerEncoderLayer( - d_model=self.latent_dim, - nhead=num_heads, - dim_feedforward=ff_size, - dropout=dropout, - activation=activation) - self.encoder = nn.TransformerEncoder(encoder_layer, - num_layers=num_layers) - - elif self.arch == "trans_dec": # Transformer Decoder - decoder_norm = nn.LayerNorm(self.latent_dim) - self.decoder = TransformerDecoder( - decoder_layer, - num_layers, - decoder_norm, - return_intermediate=return_intermediate_dec, - ) - else: - raise ValueError(f"Not supported architecture: {self.arch}!") - - def mask_cond(self, feature, mode='test', drop_prob=0.0): # train or test - bs, _, _ = feature.shape - - # classifier-free guidance - if mode == 'test': # inference - uncond_feat, con_feat = feature.chunk(2) - # con_feat = con_feat - uncond_feat = torch.zeros_like(uncond_feat) - feature = torch.cat((uncond_feat, con_feat), dim=0) - - else: # train or val mode - if drop_prob > 0.0: - mask = torch.bernoulli( - torch.ones(bs, device=feature.device) * - drop_prob).view( - bs, 1, 1) # 1-> use null_cond, 0-> use real cond - feature = feature * (1.0 - mask) - - return feature - - def get_model_kwargs( - self, - bs, - mode, - sample, - model_kwargs, - ): - - speaker_audio_encodings = model_kwargs.get('speaker_audio_encodings') - if speaker_audio_encodings is None or self.s_audio_enc_drop_prob >= 1.0: - speaker_audio_encodings = torch.zeros(size=(bs, 0, self.latent_dim)).to(sample.device) - else: - speaker_audio_encodings = self.speaker_audio_proj(speaker_audio_encodings) - speaker_audio_encodings = self.mask_cond(speaker_audio_encodings, mode, self.s_audio_enc_drop_prob) - speaker_audio_encodings = self.s_audio_scale * speaker_audio_encodings - speaker_audio_encodings = speaker_audio_encodings.permute(1, 0, 2).contiguous() - - speaker_3dmm_encodings = model_kwargs.get("speaker_3dmm_encodings") - if speaker_3dmm_encodings is None or self.s_3dmm_enc_drop_prob >= 1.0: - speaker_3dmm_encodings = torch.zeros(size=(bs, 0, self.latent_dim)).to(sample.device) - else: - speaker_3dmm_encodings = self.speaker_3dmm_proj(speaker_3dmm_encodings) - speaker_3dmm_encodings = self.mask_cond(speaker_3dmm_encodings, mode, self.s_3dmm_enc_drop_prob) - speaker_3dmm_encodings = speaker_3dmm_encodings.permute(1, 0, 2).contiguous() - - speaker_emotion_encodings = model_kwargs.get("speaker_emotion_encodings") - if speaker_emotion_encodings is None or self.s_emotion_enc_drop_prob >= 1.0: - speaker_emotion_encodings = torch.zeros(size=(bs, 0, self.latent_dim)).to(sample.device) - else: - speaker_emotion_encodings = self.emotion_proj(speaker_emotion_encodings) - speaker_emotion_encodings = self.mask_cond(speaker_emotion_encodings, mode, self.s_emotion_enc_drop_prob) - speaker_emotion_encodings = speaker_emotion_encodings.permute(1, 0, 2).contiguous() - - speaker_latent_embed = model_kwargs.get('speaker_latent_embed') - if speaker_latent_embed is None or self.s_latent_embed_drop_prob >= 1.0: - speaker_latent_embed = torch.zeros(size=(bs, 0, self.latent_dim)).to(sample.device) - else: - speaker_latent_embed = self.speaker_latent_proj(speaker_latent_embed) - speaker_latent_embed = self.mask_cond(speaker_latent_embed, mode, self.s_latent_embed_drop_prob) - speaker_latent_embed = speaker_latent_embed.permute(1, 0, 2).contiguous() - - past_listener_emotion = model_kwargs.get('past_listener_emotion') - if past_listener_emotion is None or self.past_l_emotion_drop_prob >= 1.0: - past_listener_emotion = torch.zeros(size=(bs, 0, self.latent_dim)).to(sample.device) - else: - past_listener_emotion = self.emotion_proj(past_listener_emotion) - past_listener_emotion = self.mask_cond(past_listener_emotion, mode, self.past_l_emotion_drop_prob) - past_listener_emotion = past_listener_emotion.permute(1, 0, 2).contiguous() - - return (speaker_audio_encodings, - speaker_latent_embed, - speaker_3dmm_encodings, - speaker_emotion_encodings, - past_listener_emotion) - - def _forward( - self, - sample, - time_embed, - speaker_audio_encodings, - speaker_latent_embed, - speaker_3dmm_encodings, - speaker_emotion_encodings, - past_listener_emotion, - motion_length=None, # Tensor: (bz, ) - ): - - # print("speaker_audio_encodings", speaker_audio_encodings.shape) - # print("speaker_latent_embed", speaker_latent_embed.shape) - # print("speaker_3dmm_encodings", speaker_3dmm_encodings.shape) - # print("speaker_emotion_encodings", speaker_emotion_encodings.shape) - # print("past_listener_emotion", past_listener_emotion.shape) \ - # if past_listener_emotion is not None else print("past_listener_emotion None") - - # map to latent dim - sample = self.to_emotion_embed(sample) - device = time_embed.device - - emb_latent_list = [ - time_embed, - speaker_audio_encodings, - speaker_3dmm_encodings, - speaker_emotion_encodings, - speaker_latent_embed, # TODO motion_length applied rigorously - past_listener_emotion, - ] - # [N', bs, d] - - emb_latent = torch.cat(emb_latent_list, dim=0) - - tgt_key_padding_mask = None - memory_key_padding_mask = None - - if self.task == "offline": - assert motion_length is not None, "motion_length should not be None when task is offline" - mask = lengths_to_mask(motion_length, # (bs,) - device=device, - max_len=speaker_audio_encodings.shape[0]) - # print(f"mask shape {mask.shape}") - - mask_u = torch.zeros(size=(time_embed.shape[1], - time_embed.shape[0])).bool().to(device) - # print(f"mask_u shape {mask_u.shape}") - - tgt_key_padding_mask = ~mask - memory_key_padding_mask = torch.zeros(size=(time_embed.shape[1], 0), - dtype=torch.bool).to(device) - # (bs, clip_len) - for emb in emb_latent_list: - if emb.shape[0] == mask_u.shape[-1]: - memory_key_padding_mask = torch.cat((memory_key_padding_mask, mask_u), dim=-1) - elif emb.shape[0] == mask.shape[-1]: - memory_key_padding_mask = torch.cat((memory_key_padding_mask, ~mask), dim=-1) - - else: # online task - if motion_length is not None: # during test stage - mask_l = lengths_to_mask(motion_length, - device=device, - max_len=self.window_size) - mask_s = torch.cat( - [torch.zeros_like(mask_l).bool().to(device), - ~mask_l], dim=-1) # (bz, 2 * window_size) - mask_u = torch.zeros(size=(time_embed.shape[1], - time_embed.shape[0])).bool().to(device) - - tgt_key_padding_mask = ~mask_l - memory_key_padding_mask = torch.zeros(size=(time_embed.shape[1], 0), - dtype=torch.bool).to(device) - for emb in emb_latent_list: - if emb.shape[0] == mask_u.shape[-1]: - memory_key_padding_mask = torch.cat((memory_key_padding_mask, mask_u), dim=-1) - elif emb.shape[0] == mask_s.shape[-1]: - memory_key_padding_mask = torch.cat((memory_key_padding_mask, mask_s), dim=-1) - # else: - # # memory_key_padding_mask = torch.cat((memory_key_padding_mask, mask_l), dim=-1) - # raise ValueError(f"emb.shape[0] != mask_u.shape[-1] or emb.shape[0] != mask_s.shape[-1]") - memory_key_padding_mask = None if memory_key_padding_mask.shape[-1] == 0 else memory_key_padding_mask - - if self.arch == "trans_dec": - sample = self.query_pos(sample) - emb_latent = self.mem_pos(emb_latent) - sample = self.decoder(tgt=sample, memory=emb_latent, - tgt_key_padding_mask=tgt_key_padding_mask, - memory_key_padding_mask=memory_key_padding_mask).squeeze(0) - else: - raise NotImplementedError(f"{self.arch} is not supported.") - - sample = self.to_emotion_feat(sample) - sample = sample.permute(1, 0, 2) - - return sample - - def forward_with_cond_scale( - self, - sample, # noised x_t - timesteps, - model_kwargs, - - ): - sample = torch.cat([sample] * 2, dim=0) - bs, _, _ = sample.shape - sample = sample.permute(1, 0, 2).contiguous() - - timesteps = torch.cat([timesteps] * 2, dim=0) - # with embedding permutation: [batch_size, l, encoded_dim] => [l, batch_size, encoded_dim] - time_emb = self.time_proj(timesteps) # time_embedding - time_emb = time_emb.to(dtype=sample.dtype) - time_embed = self.time_embedding(time_emb).unsqueeze(0) - - if model_kwargs is None: - model_kwargs = {} - else: - model_kwargs = model_kwargs.copy() - for k, v in model_kwargs.items(): - if model_kwargs[k] is None: - continue - model_kwargs[k] = torch.cat( - (torch.zeros_like(model_kwargs[k], dtype=model_kwargs[k].dtype), model_kwargs[k]), - dim=0) - - (speaker_audio_encodings, - speaker_latent_embed, - speaker_3dmm_encodings, - speaker_emotion_encodings, - past_listener_emotion) = ( - self.get_model_kwargs( - bs, - 'test', - sample, - model_kwargs, - ) - ) - - # motion_length = torch.cat([motion_length] * 2, dim=0) if motion_length is not None else None - prediction = self._forward( - sample, - time_embed, - speaker_audio_encodings, - speaker_latent_embed, - speaker_3dmm_encodings, - speaker_emotion_encodings, - past_listener_emotion, - model_kwargs.get('motion_length', None), - ) - - pred_uncond, pred_cond = prediction.chunk(2) - if self.guidance_scale == 1.0: # TODO to be updated - return pred_cond - - # classifier-free guidance - prediction = pred_uncond + self.guidance_scale * (pred_cond - pred_uncond) - return prediction - - def forward( - self, - sample, # noised x_t - timesteps, - model_kwargs, - **kwargs, - ): - bs, _, _ = sample.shape - sample = sample.permute(1, 0, 2).contiguous() - - # with embedding permutation: [batch_size, l, encoded_dim] => [l, batch_size, encoded_dim] - time_emb = self.time_proj(timesteps) # time_embedding - time_emb = time_emb.to(dtype=sample.dtype) - time_embed = self.time_embedding(time_emb).unsqueeze(0) - - (speaker_audio_encodings, - speaker_latent_embed, - speaker_3dmm_encodings, - speaker_emotion_encodings, - past_listener_emotion) = ( - self.get_model_kwargs( - bs, - 'train', - sample, - model_kwargs, - ) - ) - - output = self._forward( - sample, - time_embed, - speaker_audio_encodings, - speaker_latent_embed, - speaker_3dmm_encodings, - speaker_emotion_encodings, - past_listener_emotion, - model_kwargs.get('motion_length', None), - ) - - return output - - def get_model_name(self): - return self.__class__.__name__ - -# listener_latent_embed = model_kwargs.get('listener_latent_embed') -# if listener_latent_embed is None or self.l_latent_embed_drop_prob >= 1.0: -# listener_latent_embed = torch.zeros(size=(bs, 0, self.latent_dim)).to(sample.device) -# else: -# # [1, bs, encoded_dim] => [1, bs, latent_dim] -# listener_latent_embed = self.listener_latent_proj(listener_latent_embed) -# listener_latent_embed = self.mask_cond(listener_latent_embed, mode, self.l_latent_embed_drop_prob) -# listener_latent_embed = listener_latent_embed.permute(1, 0, 2).contiguous() - -# listener_personal_embed = model_kwargs.get('listener_personal_embed') -# if listener_personal_embed is None or self.l_personal_embed_drop_prob >= 1.0: -# listener_personal_embed = torch.zeros(size=(bs, 0, self.latent_dim)).to(sample.device) -# else: -# listener_personal_embed = self.listener_personal_proj(listener_personal_embed) -# listener_personal_embed = self.mask_cond(listener_personal_embed, mode, self.l_personal_embed_drop_prob) -# listener_personal_embed = listener_personal_embed.permute(1, 0, 2).contiguous() - -# elif self.condition_concat == 'feat_concat': -# assert speaker_emotion_encodings.shape[0] == speaker_3dmm_encodings.shape[0], \ -# "we need to use temporal emotion encodings or raw AU features." -# -# emb_latent = torch.cat(( -# speaker_audio_encodings, -# speaker_3dmm_encodings, -# speaker_emotion_encodings, -# ), dim=-1) # [seq_len, bs, latent_dim * 3] -# -# emb_latent = self.condition_proj(emb_latent) # (seq_len, bs, latent_dim) -# # append time embedding -# emb_latent = torch.cat((time_embed, emb_latent), dim=0) -# embed_seq_len = emb_latent.shape[0] -# -# if self.arch == "trans_enc": -# if self.concat == "concat_first": -# xseq = torch.cat((emb_latent, sample), dim=0) -# xseq = self.query_pos(xseq) -# tokens = self.encoder(xseq) -# sample = tokens[embed_seq_len:] -# elif self.concat == "concat_last": -# xseq = torch.cat((sample, emb_latent), dim=0) -# xseq = self.query_pos(xseq) -# tokens = self.encoder(xseq) -# sample = tokens[:embed_seq_len] -# else: -# raise NotImplementedError(f"{self.concat} is not supported.") -# -# elif self.arch == "trans_dec": -# sample = self.query_pos(sample) -# emb_latent = self.mem_pos(emb_latent) -# sample = self.decoder(tgt=sample, memory=emb_latent).squeeze(0) -# -# else: -# raise NotImplementedError(f"{self.arch} is not supported.") -# -# elif self.condition_concat == 'cascade': -# sample = self.query_pos(sample) -# memory = torch.cat(( -# time_embed, -# speaker_3dmm_encodings, -# speaker_latent_embed, -# ), dim=0) -# -# sample = self.transformer_fusion_ns2sv(tgt=sample, memory=memory).squeeze(0) -# memory = torch.cat((time_embed, speaker_audio_encodings), dim=0) -# sample = self.transformer_fusion_ns2sa(tgt=sample, memory=memory).squeeze(0) -# sample = self.transformer_fusion_final(tgt=sample, memory=sample).squeeze(0) -# -# # map back to original dim -# sample = self.to_emotion_feat(sample) -# sample = sample.permute(1, 0, 2) - -# if self.condition_concat == 'cascade': -# num_layer_ns2sa = 3 # noise sample <---interact---> speaker audio -# num_layer_ns2sv = 3 # noise sample <---interact---> speaker visual -# -# self.transformer_fusion_ns2sa = TransformerDecoder( -# decoder_layer, num_layer_ns2sa, nn.LayerNorm(self.latent_dim), -# return_intermediate=return_intermediate_dec, -# ) -# -# # noised sample <--> speaker visual -# self.transformer_fusion_ns2sv = TransformerDecoder( -# decoder_layer, num_layer_ns2sv, nn.LayerNorm(self.latent_dim), -# return_intermediate=return_intermediate_dec, -# ) -# -# # final interaction -# self.transformer_fusion_final = TransformerDecoder( -# decoder_layer, 1, nn.LayerNorm(self.latent_dim), -# return_intermediate=return_intermediate_dec, -# ) - -# if self.arch == "trans_enc": -# if self.concat == "concat_first": -# xseq = torch.cat((emb_latent, sample), dim=0) -# xseq = self.query_pos(xseq) -# tokens = self.encoder(xseq, src_key_padding_mask=src_key_padding_mask) -# sample = tokens[embed_seq_len:] -# elif self.concat == "concat_last": -# xseq = torch.cat((sample, emb_latent), dim=0) -# xseq = self.query_pos(xseq) -# tokens = self.encoder(xseq, src_key_padding_mask=src_key_padding_mask) -# sample = tokens[:embed_seq_len] -# else: -# raise NotImplementedError(f"{self.concat} is not supported.") diff --git a/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/transformer_denoiser_causal.py b/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/transformer_denoiser_causal.py deleted file mode 100644 index 8747a0c7e8d933f43e44475504753f8ac992e3bf..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/transformer_denoiser_causal.py +++ /dev/null @@ -1,241 +0,0 @@ -""" -CausalTransformerDenoiser -========================= -Non-invasive subclass of `TransformerDenoiser` that adds, for the OFFLINE task: - - 1. Per-timestep speaker fusion: audio_t, 3dmm_t, emotion_t -> one fused speaker - token s_t (a clean, time-aligned speaker sequence of length T instead of a - 3*T token bag). - 2. Causal lead-lag cross-attention bias: listener query at time t attends to - speaker key at time tau with an additive, per-head learnable bias b_h(t - tau) - and a hard causal mask (future speaker frames, tau > t + lookahead, are -inf). - Different heads specialise in different reaction delays. - 3. Coarse-to-fine conditioning: a causal GRU over the fused speaker sequence - predicts the listener's 8-class facial-expression "plan" per timestep - (`_coarse_logits`, surfaced for an explicit CE loss) and FiLM-modulates the - listener tokens before the decoder. - -Everything else (diffusion wrapper, EEG head, loss plumbing for the other terms, -checkpointing) is inherited unchanged. The coarse logits are stashed on -`self._coarse_logits` so the matcher can surface them to the loss without -touching the model's tensor return signature. -""" -from typing import Optional - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from framework.motion_diffusion.diffusion.diffusion_decoder.transformer_denoiser import ( - TransformerDenoiser, - lengths_to_mask, -) - - -class CausalTransformerDenoiser(TransformerDenoiser): - def __init__( - self, - num_heads: int = 4, - lag_max: int = 60, - lag_lookahead: int = 0, - coarse_classes: int = 8, - coarse_hidden: int = 256, - coarse_emo_start: int = 17, - use_lag_bias: bool = True, - use_coarse: bool = True, - **kwargs, - ) -> None: - super().__init__(num_heads=num_heads, **kwargs) - - self.num_heads = num_heads - self.lag_max = int(lag_max) - self.lag_lookahead = int(lag_lookahead) - self.coarse_classes = int(coarse_classes) - self.coarse_emo_start = int(coarse_emo_start) - self.use_lag_bias = bool(use_lag_bias) - self.use_coarse = bool(use_coarse) - - d = self.latent_dim - - # (1) per-timestep speaker fusion: [audio | 3dmm | emotion] (3d) -> d - self.fuse_proj = nn.Linear(3 * d, d) - - # (2) causal lead-lag attention bias table (per head). - # buckets: [-lookahead .. lag_max] exact, plus one "far past" bucket. - self.num_lag_buckets = self.lag_lookahead + self.lag_max + 1 + 1 - self.lag_bias = nn.Parameter(torch.zeros(num_heads, self.num_lag_buckets)) - - # (3) coarse-to-fine head: causal GRU over fused speaker -> per-t plan. - self.coarse_gru = nn.GRU(d, coarse_hidden, num_layers=1, batch_first=False) - self.coarse_out = nn.Linear(coarse_hidden, self.coarse_classes) - self.coarse_film = nn.Linear(coarse_hidden, 2 * d) - # start FiLM near identity (gamma~0, beta~0) - nn.init.zeros_(self.coarse_film.weight) - nn.init.zeros_(self.coarse_film.bias) - - self._coarse_logits: Optional[torch.Tensor] = None - - # (P4) optional per-person FiLM over the coarse GRU hidden state. Set - # externally (by the personalized modifier) to a (gamma, beta) tuple of - # broadcastable tensors right before the forward; None => generic plan. - self._person_coarse_film = None - - # ------------------------------------------------------------------ utils - def _fuse_speaker(self, audio, mm3d, emotion, T, bs, device, dtype): - """Fuse the three per-frame speaker encodings into one token / frame. - - Each input is (Ts, bs, d) or length-0 (dropped). Missing channels are - replaced by zeros so the fusion projection always sees 3*d. - """ - def _fix(x): - if x is None or x.shape[0] == 0: - return torch.zeros(T, bs, self.latent_dim, device=device, dtype=dtype) - return x - - audio, mm3d, emotion = _fix(audio), _fix(mm3d), _fix(emotion) - fused = self.fuse_proj(torch.cat([audio, mm3d, emotion], dim=-1)) # (T, bs, d) - return fused - - def _build_lag_mask(self, T, Ts, n_global, bs, device, dtype, offset=0): - """Additive cross-attention mask (bs*num_heads, T, n_global + Ts). - - Global tokens (time / latent / past) get bias 0 and are always visible. - Speaker columns get per-head lag bias b_h(t-tau); future (tau > t + - lookahead) is -inf (hard causality). - """ - # `offset` aligns the listener window to the speaker window in absolute - # frame index: listener frame t corresponds to speaker frame (t + offset). - # offline: Ts == T, offset == 0. online: speaker leads by offset = Ts - T. - t_idx = (torch.arange(T, device=device) + offset).unsqueeze(1) # (T, 1) - tau = torch.arange(Ts, device=device).unsqueeze(0) # (1, Ts) - delta = t_idx - tau # (T, Ts) - - future = delta < (-self.lag_lookahead) # speaker in the future - far = delta > self.lag_max # very old speaker - bucket = (delta.clamp(min=-self.lag_lookahead, max=self.lag_max) - + self.lag_lookahead) # [0 .. lookahead+lag_max] - bucket = bucket.masked_fill(far, self.num_lag_buckets - 1) - bucket = bucket.clamp(0, self.num_lag_buckets - 1).long() - - bias = self.lag_bias[:, bucket] # (num_heads, T, Ts) - bias = bias.masked_fill(future.unsqueeze(0), float("-inf")) - - gbias = torch.zeros(self.num_heads, T, n_global, device=device, dtype=bias.dtype) - full = torch.cat([gbias, bias], dim=2) # (num_heads, T, S_mem) - full = full.unsqueeze(0).expand(bs, -1, -1, -1).reshape( - bs * self.num_heads, T, n_global + Ts) - return full.to(dtype) - - # --------------------------------------------------------------- forward - def _forward( - self, - sample, - time_embed, - speaker_audio_encodings, - speaker_latent_embed, - speaker_3dmm_encodings, - speaker_emotion_encodings, - past_listener_emotion, - motion_length=None, - ): - # Causal path supports offline (full sequence, Ts == T) and online - # (windowed: speaker s_w leads listener l_w by offset = Ts - T frames). - if self.arch != "trans_dec" or self.task not in ("offline", "online"): - return super()._forward( - sample, time_embed, speaker_audio_encodings, speaker_latent_embed, - speaker_3dmm_encodings, speaker_emotion_encodings, - past_listener_emotion, motion_length) - - device = time_embed.device - dtype = sample.dtype - - sample = self.to_emotion_embed(sample) # (T, bs, d) - T, bs, _ = sample.shape - - # (1) fuse speaker channels into a time-aligned sequence ---------------- - Ts = speaker_audio_encodings.shape[0] - if Ts == 0: - Ts = max(speaker_3dmm_encodings.shape[0], speaker_emotion_encodings.shape[0], T) - fused = self._fuse_speaker( - speaker_audio_encodings, speaker_3dmm_encodings, speaker_emotion_encodings, - Ts, bs, device, dtype) # (Ts, bs, d) - offset = Ts - T # speaker leads listener (online: 30, offline: 0) - - # (3) coarse-to-fine plan (causal GRU over fused speaker) --------------- - # The GRU runs over all speaker frames; the listener-aligned hidden states - # are the last T (indices [offset : offset+T]) -> one plan per listener frame. - if self.use_coarse: - h, _ = self.coarse_gru(fused) # (Ts, bs, H) - h_l = h[offset:offset + T] # (T, bs, H) aligned to listener - # (P4) personalize the plan: FiLM the hidden state with the person's - # embedding so BOTH the 8-class logits and the downstream fine FiLM - # reflect "how this person reacts" (zero-init => identity at start). - pcf = self._person_coarse_film - if pcf is not None: - p_gamma, p_beta = pcf # each (·, H) broadcastable - if p_gamma.shape[0] == 1 and bs != 1: - p_gamma = p_gamma.expand(bs, -1) - p_beta = p_beta.expand(bs, -1) - elif p_gamma.shape[0] != bs: - raise ValueError( - f"Personal coarse batch {p_gamma.shape[0]} cannot broadcast " - f"to denoiser batch {bs}" - ) - h_l = (1.0 + p_gamma.view(1, bs, -1)) * h_l + p_beta.view(1, bs, -1) - coarse_logits = self.coarse_out(h_l) # (T, bs, C) - self._coarse_logits = coarse_logits.permute(1, 0, 2).contiguous() # (bs, T, C) - film = self.coarse_film(h_l) # (T, bs, 2d) - gamma, beta = film.chunk(2, dim=-1) # (T, bs, d) each - else: - self._coarse_logits = None - gamma = beta = None - - # listener tokens + positional encoding + FiLM conditioning ------------ - sample = self.query_pos(sample) - if self.use_coarse and gamma is not None: - sample = (1.0 + gamma) * sample + beta - - # (2) assemble memory: [global tokens ... | fused speaker seq] ---------- - # globals (time / speaker-latent / past-listener) are always visible - # (past-listener frames all precede the target window -> no causal mask). - global_tokens = [time_embed] - if speaker_latent_embed is not None and speaker_latent_embed.shape[0] > 0: - global_tokens.append(speaker_latent_embed) - if past_listener_emotion is not None and past_listener_emotion.shape[0] > 0: - global_tokens.append(past_listener_emotion) - n_global = sum(g.shape[0] for g in global_tokens) - - memory = torch.cat(global_tokens + [fused], dim=0) # (n_global + Ts, bs, d) - memory = self.mem_pos(memory) - - # padding masks -------------------------------------------------------- - tgt_key_padding_mask = None - memory_key_padding_mask = None - if motion_length is not None: - l_valid = lengths_to_mask(motion_length, device=device, max_len=T) # (bs, T) - # speaker validity: the leading `offset` context frames are always - # valid; the last T speaker frames share the listener's validity. - if offset > 0: - s_valid = torch.cat( - [torch.ones(bs, offset, dtype=torch.bool, device=device), l_valid], dim=1) - else: - s_valid = l_valid # (bs, Ts==T) - tgt_key_padding_mask = ~l_valid - g_pad = torch.zeros(bs, n_global, dtype=torch.bool, device=device) # globals never padded - memory_key_padding_mask = torch.cat([g_pad, ~s_valid], dim=1) # (bs, n_global+Ts) - - # (2) causal lead-lag additive mask ----------------------------------- - memory_mask = None - if self.use_lag_bias: - memory_mask = self._build_lag_mask(T, Ts, n_global, bs, device, dtype, offset=offset) - - sample = self.decoder( - tgt=sample, memory=memory, - memory_mask=memory_mask, - tgt_key_padding_mask=tgt_key_padding_mask, - memory_key_padding_mask=memory_key_padding_mask, - ).squeeze(0) - - sample = self.to_emotion_feat(sample) - sample = sample.permute(1, 0, 2) - return sample diff --git a/personalised/code/framework/motion_diffusion/diffusion/diffusion_prior/__init__.py b/personalised/code/framework/motion_diffusion/diffusion/diffusion_prior/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/personalised/code/framework/motion_diffusion/diffusion/diffusion_prior/rotary_embedding_torch.py b/personalised/code/framework/motion_diffusion/diffusion/diffusion_prior/rotary_embedding_torch.py deleted file mode 100644 index 13746c145860e110d562164a5e2ee315f023a738..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/diffusion_prior/rotary_embedding_torch.py +++ /dev/null @@ -1,282 +0,0 @@ -from math import pi, log - -import torch -from torch.nn import Module, ModuleList -from torch.cuda.amp import autocast -from torch import nn, einsum, broadcast_tensors, Tensor - -from einops import rearrange, repeat - -from beartype import beartype -from beartype.typing import Literal, Union, Optional - -# helper functions - -def exists(val): - return val is not None - -def default(val, d): - return val if exists(val) else d - -# broadcat, as tortoise-tts was using it - -def broadcat(tensors, dim = -1): - broadcasted_tensors = broadcast_tensors(*tensors) - return torch.cat(broadcasted_tensors, dim = dim) - -# rotary embedding helper functions - -def rotate_half(x): - x = rearrange(x, '... (d r) -> ... d r', r = 2) - x1, x2 = x.unbind(dim = -1) - x = torch.stack((-x2, x1), dim = -1) - return rearrange(x, '... d r -> ... (d r)') - -@autocast(enabled = False) -def apply_rotary_emb(freqs, t, start_index = 0, scale = 1., seq_dim = -2): - if t.ndim == 3: - seq_len = t.shape[seq_dim] - freqs = freqs[-seq_len:].to(t) - - rot_dim = freqs.shape[-1] - end_index = start_index + rot_dim - - assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}' - - t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:] - t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale) - return torch.cat((t_left, t, t_right), dim = -1) - -# learned rotation helpers - -def apply_learned_rotations(rotations, t, start_index = 0, freq_ranges = None): - if exists(freq_ranges): - rotations = einsum('..., f -> ... f', rotations, freq_ranges) - rotations = rearrange(rotations, '... r f -> ... (r f)') - - rotations = repeat(rotations, '... n -> ... (n r)', r = 2) - return apply_rotary_emb(rotations, t, start_index = start_index) - -# classes - -class RotaryEmbedding(Module): - @beartype - def __init__( - self, - dim, - custom_freqs: Optional[Tensor] = None, - freqs_for: Union[ - Literal['lang'], - Literal['pixel'], - Literal['constant'] - ] = 'lang', - theta = 10000, - max_freq = 10, - num_freqs = 1, - learned_freq = False, - use_xpos = False, - xpos_scale_base = 512, - interpolate_factor = 1., - theta_rescale_factor = 1., - seq_before_head_dim = False, - cache_if_possible = True - ): - super().__init__() - # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning - # has some connection to NTK literature - # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/ - - theta *= theta_rescale_factor ** (dim / (dim - 2)) - - self.freqs_for = freqs_for - - if exists(custom_freqs): - freqs = custom_freqs - elif freqs_for == 'lang': - freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) - elif freqs_for == 'pixel': - freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi - elif freqs_for == 'constant': - freqs = torch.ones(num_freqs).float() - - self.cache_if_possible = cache_if_possible - - self.tmp_store('cached_freqs', None) - self.tmp_store('cached_scales', None) - - self.freqs = nn.Parameter(freqs, requires_grad = learned_freq) - - self.learned_freq = learned_freq - - # dummy for device - - self.tmp_store('dummy', torch.tensor(0)) - - # default sequence dimension - - self.seq_before_head_dim = seq_before_head_dim - self.default_seq_dim = -3 if seq_before_head_dim else -2 - - # interpolation factors - - assert interpolate_factor >= 1. - self.interpolate_factor = interpolate_factor - - # xpos - - self.use_xpos = use_xpos - if not use_xpos: - self.tmp_store('scale', None) - return - - scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim) - self.scale_base = xpos_scale_base - self.tmp_store('scale', scale) - - @property - def device(self): - return self.dummy.device - - def tmp_store(self, key, value): - self.register_buffer(key, value, persistent = False) - - def get_seq_pos(self, seq_len, device, dtype, offset = 0): - return (torch.arange(seq_len, device = device, dtype = dtype) + offset) / self.interpolate_factor - - def rotate_queries_or_keys(self, t, seq_dim = None, offset = 0, freq_seq_len = None): - seq_dim = default(seq_dim, self.default_seq_dim) - - assert not self.use_xpos, 'you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings' - - device, dtype, seq_len = t.device, t.dtype, t.shape[seq_dim] - - if exists(freq_seq_len): - assert freq_seq_len >= seq_len - seq_len = freq_seq_len - - freqs = self.forward(self.get_seq_pos(seq_len, device = device, dtype = dtype, offset = offset), seq_len = seq_len, offset = offset) - - if seq_dim == -3: - freqs = rearrange(freqs, 'n d -> n 1 d') - - return apply_rotary_emb(freqs, t, seq_dim = seq_dim) - - def rotate_queries_with_cached_keys(self, q, k, seq_dim = None, offset = 0): - seq_dim = default(seq_dim, self.default_seq_dim) - - q_len, k_len = q.shape[seq_dim], k.shape[seq_dim] - assert q_len <= k_len - rotated_q = self.rotate_queries_or_keys(q, seq_dim = seq_dim, freq_seq_len = k_len) - rotated_k = self.rotate_queries_or_keys(k, seq_dim = seq_dim) - - rotated_q = rotated_q.type(q.dtype) - rotated_k = rotated_k.type(k.dtype) - - return rotated_q, rotated_k - - def rotate_queries_and_keys(self, q, k, seq_dim = None): - seq_dim = default(seq_dim, self.default_seq_dim) - - assert self.use_xpos - device, dtype, seq_len = q.device, q.dtype, q.shape[seq_dim] - - seq = self.get_seq_pos(seq_len, dtype = dtype, device = device) - - freqs = self.forward(seq, seq_len = seq_len) - scale = self.get_scale(seq, seq_len = seq_len).to(dtype) - - if seq_dim == -3: - freqs = rearrange(freqs, 'n d -> n 1 d') - scale = rearrange(scale, 'n d -> n 1 d') - - rotated_q = apply_rotary_emb(freqs, q, scale = scale, seq_dim = seq_dim) - rotated_k = apply_rotary_emb(freqs, k, scale = scale ** -1, seq_dim = seq_dim) - - rotated_q = rotated_q.type(q.dtype) - rotated_k = rotated_k.type(k.dtype) - - return rotated_q, rotated_k - - @beartype - def get_scale( - self, - t: Tensor, - seq_len: Optional[int] = None, - offset = 0 - ): - assert self.use_xpos - - should_cache = ( - self.cache_if_possible and - exists(seq_len) - ) - - if ( - should_cache and \ - exists(self.cached_scales) and \ - (seq_len + offset) <= self.cached_scales.shape[0] - ): - return self.cached_scales[offset:(offset + seq_len)] - - scale = 1. - if self.use_xpos: - power = (t - len(t) // 2) / self.scale_base - scale = self.scale ** rearrange(power, 'n -> n 1') - scale = torch.cat((scale, scale), dim = -1) - - if should_cache: - self.tmp_store('cached_scales', scale) - - return scale - - def get_axial_freqs(self, *dims): - Colon = slice(None) - all_freqs = [] - - for ind, dim in enumerate(dims): - if self.freqs_for == 'pixel': - pos = torch.linspace(-1, 1, steps = dim, device = self.device) - else: - pos = torch.arange(dim, device = self.device) - - freqs = self.forward(pos, seq_len = dim) - - all_axis = [None] * len(dims) - all_axis[ind] = Colon - - new_axis_slice = (Ellipsis, *all_axis, Colon) - all_freqs.append(freqs[new_axis_slice]) - - all_freqs = broadcast_tensors(*all_freqs) - return torch.cat(all_freqs, dim = -1) - - @autocast(enabled = False) - def forward( - self, - t: Tensor, - seq_len = None, - offset = 0 - ): - should_cache = ( - self.cache_if_possible and \ - not self.learned_freq and \ - exists(seq_len) and \ - self.freqs_for != 'pixel' - ) - - if ( - should_cache and \ - exists(self.cached_freqs) and \ - (offset + seq_len) <= self.cached_freqs.shape[0] - ): - return self.cached_freqs[offset:(offset + seq_len)].detach() - - freqs = self.freqs - - freqs = einsum('..., f -> ... f', t.type(freqs.dtype), freqs) - freqs = repeat(freqs, '... n -> ... (n r)', r = 2) - - if should_cache: - self.tmp_store('cached_freqs', freqs.detach()) - - return freqs diff --git a/personalised/code/framework/motion_diffusion/diffusion/diffusion_prior/transformer_prior.py b/personalised/code/framework/motion_diffusion/diffusion/diffusion_prior/transformer_prior.py deleted file mode 100644 index 7cee0e0159288a52609755f203fe88bacc3c553b..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/diffusion_prior/transformer_prior.py +++ /dev/null @@ -1,435 +0,0 @@ -import torch -import torch.nn as nn -from torch import einsum -import torch.nn.functional as F -from einops import rearrange, repeat -import math -from framework.motion_diffusion.diffusion.diffusion_prior.rotary_embedding_torch import RotaryEmbedding -from enum import Enum -from einops.layers.torch import Rearrange -from framework.motion_diffusion.diffusion.utils.util import prob_mask_like - - -def exists(val): - return val is not None - - -def l2norm(t): - return F.normalize(t, dim=-1) - - -class Activation(Enum): - none = 'none' - relu = 'relu' - lrelu = 'lrelu' - silu = 'silu' - tanh = 'tanh' - - def get_act(self): - if self == Activation.none: - return nn.Identity() - elif self == Activation.relu: - return nn.ReLU() - elif self == Activation.lrelu: - return nn.LeakyReLU(negative_slope=0.2) - elif self == Activation.silu: - return nn.SiLU() - elif self == Activation.tanh: - return nn.Tanh() - else: - raise NotImplementedError() - - -class SwiGLU(nn.Module): - """ used successfully in https://arxiv.org/abs/2204.0231 """ - - def forward(self, x): - x, gate = x.chunk(2, dim=-1) - return x * F.silu(gate) - - -def timestep_embedding(timesteps, dim, max_period=10000, dtype=torch.float32): - """ - Create sinusoidal timestep embeddings. - :param timesteps: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param dim: the dimension of the output. - :param max_period: controls the minimum frequency of the embeddings. - :return: an [N x dim] Tensor of positional embeddings. - """ - half = dim // 2 - freqs = torch.exp( - -math.log(max_period) * torch.arange(start=0, end=half, dtype=dtype) / half - ).to(device=timesteps.device) - args = timesteps[:, None].type(dtype) * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) - return embedding - - -class RelPosBias(nn.Module): - def __init__( - self, - heads=8, - num_buckets=32, - max_distance=128, - ): - super().__init__() - self.num_buckets = num_buckets - self.max_distance = max_distance - self.relative_attention_bias = nn.Embedding(num_buckets, heads) - - @staticmethod - def _relative_position_bucket( - relative_position, - num_buckets=32, - max_distance=128 - ): - n = -relative_position - n = torch.max(n, torch.zeros_like(n)) - - max_exact = num_buckets // 2 - is_small = n < max_exact - - val_if_large = max_exact + (torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * ( - num_buckets - max_exact)).long() - val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1)) - return torch.where(is_small, n, val_if_large) - - def forward(self, i, j, *, device): - q_pos = torch.arange(i, dtype=torch.long, device=device) - k_pos = torch.arange(j, dtype=torch.long, device=device) - rel_pos = rearrange(k_pos, 'j -> 1 j') - rearrange(q_pos, 'i -> i 1') - rp_bucket = self._relative_position_bucket(rel_pos, num_buckets=self.num_buckets, - max_distance=self.max_distance) - values = self.relative_attention_bias(rp_bucket) - return rearrange(values, 'i j h -> h i j') - - -class LayerNorm(nn.Module): - def __init__(self, dim, eps=1e-5, fp16_eps=1e-3, stable=False): - super().__init__() - self.eps = eps - self.fp16_eps = fp16_eps - self.stable = stable - self.g = nn.Parameter(torch.ones(dim)) - - def forward(self, x): - eps = self.eps if x.dtype == torch.float32 else self.fp16_eps - - if self.stable: - x = x / x.amax(dim=-1, keepdim=True).detach() - - var = torch.var(x, dim=-1, unbiased=False, keepdim=True) - mean = torch.mean(x, dim=-1, keepdim=True) - return (x - mean) * (var + eps).rsqrt() * self.g - - -def FeedForward( - dim, - mult=4, - dropout=0., - post_activation_norm=False -): - """ post-activation norm https://arxiv.org/abs/2110.09456 """ - - inner_dim = int(mult * dim) - return nn.Sequential( - LayerNorm(dim), - nn.Linear(dim, inner_dim * 2, bias=False), - SwiGLU(), - LayerNorm(inner_dim) if post_activation_norm else nn.Identity(), - nn.Dropout(dropout), - nn.Linear(inner_dim, dim, bias=False) - ) - - -class Attention(nn.Module): - def __init__( - self, - dim, - *, - dim_head=64, - heads=8, - dropout=0., - causal=False, - rotary_emb=None, - cosine_sim=True, - cosine_sim_scale=16 - ): - super().__init__() - self.scale = cosine_sim_scale if cosine_sim else (dim_head ** -0.5) - self.cosine_sim = cosine_sim - - self.heads = heads - inner_dim = dim_head * heads - - self.causal = causal - self.norm = LayerNorm(dim) - self.dropout = nn.Dropout(dropout) - - self.null_kv = nn.Parameter(torch.randn(2, dim_head)) - self.to_q = nn.Linear(dim, inner_dim, bias=False) - self.to_kv = nn.Linear(dim, dim_head * 2, bias=False) - - self.rotary_emb = rotary_emb - - self.to_out = nn.Sequential( - nn.Linear(inner_dim, dim, bias=False), - LayerNorm(dim) - ) - - def forward(self, x, mask=None, attn_bias=None): - b, n, device = *x.shape[:2], x.device - - x = self.norm(x) - q, k, v = (self.to_q(x), *self.to_kv(x).chunk(2, dim=-1)) - - q = rearrange(q, 'b n (h d) -> b h n d', h=self.heads) - q = q * self.scale - - # rotary embeddings - if exists(self.rotary_emb): - q, k = map(self.rotary_emb.rotate_queries_or_keys, (q, k)) - - # add null key / value for classifier free guidance in prior net - nk, nv = map(lambda t: repeat(t, 'd -> b 1 d', b=b), self.null_kv.unbind(dim=-2)) - k = torch.cat((nk, k), dim=-2) - v = torch.cat((nv, v), dim=-2) - - # whether to use cosine sim - if self.cosine_sim: - q, k = map(l2norm, (q, k)) - - q, k = map(lambda t: t * math.sqrt(self.scale), (q, k)) - - # calculate query / key similarities - sim = einsum('b h i d, b j d -> b h i j', q, k) - - # relative positional encoding (T5 style) - if exists(attn_bias): - sim = sim + attn_bias - - # masking - max_neg_value = -torch.finfo(sim.dtype).max - - if exists(mask): - mask = F.pad(mask, (1, 0), value=True) - mask = rearrange(mask, 'b j -> b 1 1 j') - sim = sim.masked_fill(~mask, max_neg_value) - - if self.causal: - i, j = sim.shape[-2:] - causal_mask = torch.ones((i, j), dtype=torch.bool, device=device).triu(j - i + 1) - sim = sim.masked_fill(causal_mask, max_neg_value) - - # attention - attn = sim.softmax(dim=-1, dtype=torch.float32) - attn = attn.type(sim.dtype) - attn = self.dropout(attn) - - # aggregate values - out = einsum('b h i j, b j d -> b h i d', attn, v) - - out = rearrange(out, 'b h n d -> b n (h d)') - return self.to_out(out) - - -class CausalTransformer(nn.Module): - def __init__( - self, - *, - dim=512, # latent_dim - depth=4, - dim_head=64, - heads=8, - ff_mult=4, - norm_in=False, - norm_out=True, - attn_dropout=0., - ff_dropout=0., - final_proj=True, - normformer=False, - rotary_emb=True, - ): - super().__init__() - self.init_norm = LayerNorm(dim) if norm_in else nn.Identity() # from latest BLOOM model and Yandex's YaLM - - self.rel_pos_bias = RelPosBias(heads=heads) - - rotary_emb = RotaryEmbedding(dim=min(32, dim_head)) if rotary_emb else None - - self.layers = nn.ModuleList([]) - for _ in range(depth): - self.layers.append(nn.ModuleList([ - Attention(dim=dim, causal=True, dim_head=dim_head, heads=heads, dropout=attn_dropout, - rotary_emb=rotary_emb), - FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout, post_activation_norm=normformer) - ])) - - self.norm = LayerNorm(dim, - stable=True) if norm_out else nn.Identity() - self.project_out = nn.Linear(dim, dim, bias=False) if final_proj else nn.Identity() - - def forward(self, x): - n, device = x.shape[1], x.device - - x = self.init_norm(x) - - attn_bias = self.rel_pos_bias(n, n + 1, device=device) - - for attn, ff in self.layers: - x = attn(x, attn_bias=attn_bias) + x - x = ff(x) + x - - out = self.norm(x) - return self.project_out(out) - - -class DiffusionPriorNetwork(nn.Module): - def __init__(self, - audio_dim=78, - window_size=50, - _3dmm_dim=58, - speaker_emb_dim=512, - latent_dim=512, - depth=4, - num_time_layers=2, - num_time_embeds=1, - num_time_emb_channels=64, - time_last_act=False, - activation=Activation.silu, - use_learned_query=True, - s_audio_cond_drop_prob=0.0, - s_latentemb_cond_drop_prob=0.0, - s_3dmm_cond_drop_prob=0.0, - guidance_scale=1.0, - **kwargs): - super().__init__() - - self.window_size = window_size - self.latent_dim = latent_dim - self.use_learned_query = use_learned_query - self.num_time_emb_channels = num_time_emb_channels - - layers = [] - for i in range(num_time_layers): - if i == 0: - a = num_time_emb_channels - b = latent_dim - else: - a = latent_dim - b = latent_dim - layers.append(nn.Linear(a, b)) - if i < num_time_layers - 1 or time_last_act: - layers.append(activation.get_act()) - layers.append(Rearrange('b (n d) -> b n d', n=num_time_embeds)) - self.time_embed = nn.Sequential(*layers) - - self.to_audio_encodings = nn.Linear(audio_dim, latent_dim) if audio_dim != latent_dim else nn.Identity() - self.to_speaker_latentemb = nn.Linear(speaker_emb_dim, latent_dim) \ - if speaker_emb_dim != latent_dim else nn.Identity() - self.to_speaker_3dmmenc = nn.Linear(_3dmm_dim, latent_dim) if _3dmm_dim != latent_dim else nn.Identity() - - self.learned_query = nn.Parameter(torch.randn(latent_dim)) - - self.guidance_scale = guidance_scale - self.s_audio_cond_drop_prob = s_audio_cond_drop_prob - self.s_latentemb_cond_drop_prob = s_latentemb_cond_drop_prob - self.s_3dmm_cond_drop_prob = s_3dmm_cond_drop_prob - self.null_s_audio_encodings = nn.Parameter(torch.randn(size=(1, self.window_size, self.latent_dim))) - self.null_s_latent_embed = nn.Parameter(torch.randn(size=(1, 1, self.latent_dim))) - self.null_s_3dmm_encodings = nn.Parameter(torch.randn(size=(1, self.window_size, self.latent_dim))) - - self.causal_transformer = CausalTransformer(dim=latent_dim, depth=depth, **kwargs) - - def forward_with_cond_scale(self, x_t, t, model_kwargs): - self.s_audio_cond_drop_prob = 0.0 - self.s_latentemb_cond_drop_prob = 0.0 - self.s_3dmm_cond_drop_prob = 0.0 - logits = self.forward(x_t, t, model_kwargs) - - if self.guidance_scale <= 1.0: - return logits - - self.s_audio_cond_drop_prob = 1.0 - self.s_latentemb_cond_drop_prob = 1.0 - self.s_3dmm_cond_drop_prob = 1.0 - null_logits = self.forward(x_t, t, model_kwargs) - - return null_logits + (logits - null_logits) * self.guidance_scale - - def forward(self, x_t, t, model_kwargs): - assert x_t.shape[2] == self.latent_dim, \ - "x_t should have the same dimension as the latent dimension." - - t = timestep_embedding(t, self.num_time_emb_channels) - time_emb = self.time_embed(t) - bs, _, _ = time_emb.shape - - s_audio_encodings = model_kwargs.get("speaker_audio_encodings", - torch.zeros(size=(bs, 0, self.latent_dim)).to(x_t.device)) - if s_audio_encodings.shape[1] > 0: - s_audio_encodings = self.to_audio_encodings(s_audio_encodings) # mapping - - audio_keep_mask = ( - prob_mask_like((s_audio_encodings.shape[0],), - (1 - self.s_audio_cond_drop_prob), device=x_t.device)) - audio_keep_mask = rearrange(audio_keep_mask, 'b -> b 1 1') - s_audio_encodings = torch.where( - audio_keep_mask, - s_audio_encodings, - self.null_s_audio_encodings.to(s_audio_encodings.device) - ) - speaker_audio_encodings = s_audio_encodings - - s_latent_embed = model_kwargs.get("speaker_latent_emb", - torch.zeros(size=(bs, 0, self.latent_dim)).to(x_t.device)) - if s_latent_embed.shape[1] > 0: - s_latent_embed = self.to_speaker_latentemb(s_latent_embed) - latentemb_keep_mask = ( - prob_mask_like((s_latent_embed.shape[0],), - (1 - self.s_latentemb_cond_drop_prob), device=s_latent_embed.device)) - latentemb_keep_mask = rearrange(latentemb_keep_mask, 'b -> b 1 1') - s_latent_embed = torch.where( - latentemb_keep_mask, - s_latent_embed, - self.null_s_latent_embed.to(s_latent_embed.device) - ) - speaker_latent_emb = s_latent_embed - - s_3dmm_encodings = model_kwargs.get("speaker_3dmm_encodings", - torch.zeros(size=(bs, 0, self.latent_dim)).to(x_t.device)) - if s_3dmm_encodings.shape[1] > 0: - s_3dmm_encodings = self.to_speaker_3dmmenc(s_3dmm_encodings) - _3dmm_keep_mask = ( - prob_mask_like((s_3dmm_encodings.shape[0],), - (1 - self.s_3dmm_cond_drop_prob), device=s_3dmm_encodings.device)) - _3dmm_keep_mask = rearrange(_3dmm_keep_mask, 'b -> b 1 1') - s_3dmm_encodings = torch.where( - _3dmm_keep_mask, - s_3dmm_encodings, - self.null_s_3dmm_encodings.to(s_3dmm_encodings.device) - ) - speaker_3dmm_encodings = s_3dmm_encodings - - learned_queries = repeat(self.learned_query, 'd -> b 1 d', b=bs) \ - if self.use_learned_query else torch.zeros(size=(bs, 0, self.latent_dim)).to(x_t.device) - - tokens = torch.cat(( - speaker_audio_encodings, # shape: (bs, window_size, dim) - speaker_latent_emb, # shape: (bs, 1, dim) - speaker_3dmm_encodings, # shape: (bs, window_size, dim) - time_emb, # shape: (bs, 1, dim) - x_t, # shape: (bs, 1, dim) - learned_queries, # shape: (bs, 1, dim) - ), dim=-2) # (bs, N', dim) - - tokens = self.causal_transformer(tokens) - pred_image_embed = tokens[:, -1:, :] - - return pred_image_embed - - def get_model_name(self): - return self.__class__.__name__ diff --git a/personalised/code/framework/motion_diffusion/diffusion/gaussian_diffusion.py b/personalised/code/framework/motion_diffusion/diffusion/gaussian_diffusion.py deleted file mode 100644 index 61c8276570aba103b10403e505135d1cb3789cb4..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/gaussian_diffusion.py +++ /dev/null @@ -1,1161 +0,0 @@ -""" -This code started out as a PyTorch port of Ho et al's diffusion modules: -https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py - -Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules. - -Code adapted from: -https://github.com/BarqueroGerman/BeLFusion -""" - -import enum -import math -from cv2 import CAP_PROP_XI_AUTO_BANDWIDTH_CALCULATION -import numpy as np -import torch as th -import torch.nn as nn -from einops import rearrange -from framework.motion_diffusion.diffusion.utils.util import prob_mask_like - -LOSSES_TYPES = ["mse", "mse_l1", ] -MSE, MSE_L1 = LOSSES_TYPES - - -def mean_flat(tensor): - """ - Take the mean over all non-batch dimensions. - """ - return tensor.mean(dim=list(range(1, len(tensor.shape)))) - - -def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): - """ - Get a pre-defined beta schedule for the given name. - - The beta schedule library consists of beta schedules which remain similar - in the limit of num_diffusion_timesteps. - Beta schedules may be added, but should not be removed or changed once - they are committed to maintain backwards compatibility. - """ - if schedule_name == "linear": - # Linear schedule from Ho et al, extended to work for any number of - # diffusion steps. - scale = 1000 / num_diffusion_timesteps - beta_start = scale * 0.0001 - beta_end = scale * 0.02 - return np.linspace( - beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64 - ) - elif schedule_name == "cosine": - return betas_for_alpha_bar( - num_diffusion_timesteps, - lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, - ) - elif 'sqrt' in schedule_name: - sqrt_schedulers = { - "10sqrt1e-4": lambda t: max(0, 1 - np.power(t + 0.0001, 1 / 10)), - "5sqrt1e-4": lambda t: max(0, 1 - np.power(t + 0.0001, 1 / 5)), - "3sqrt1e-4": lambda t: max(0, 1 - np.power(t + 0.0001, 1 / 3)), - "sqrt1e-4": lambda t: max(0, 1 - np.sqrt(t + 0.0001)), - "sqrt2e-2": lambda t: max(0, 1 - np.sqrt(t + 0.02)), - "sqrt5e-2": lambda t: max(0, 1 - np.sqrt(t + 0.05)), - "sqrt1e-1": lambda t: max(0, 1 - np.sqrt(t + 0.1)), - "sqrt2e-2": lambda t: max(0, 1 - np.sqrt(t + 0.2)), - } - assert schedule_name in sqrt_schedulers.keys(), f"Unknown sqrt scheduler {schedule_name}" - return betas_for_alpha_bar( - num_diffusion_timesteps, - sqrt_schedulers[schedule_name], - ) - else: - raise NotImplementedError(f"unknown beta schedule: {schedule_name}") - - -def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, - which defines the cumulative product of (1-beta) over time from t = [0,1]. - - :param num_diffusion_timesteps: the number of betas to produce. - :param alpha_bar: a lambda that takes an argument t from 0 to 1 and - produces the cumulative product of (1-beta) up to that - part of the diffusion process. - :param max_beta: the maximum beta to use; use values lower than 1 to - prevent singularities. - """ - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append( - min(1 - alpha_bar(t2) / max(0.00001, alpha_bar(t1)), max_beta)) # the max is to prevent singularities - return np.array(betas) - - -class ModelMeanType(enum.Enum): - """ - Which type of output the model predicts. - """ - - PREVIOUS_X = enum.auto() # the model predicts x_{t-1} - START_X = enum.auto() # the model predicts x_0 - EPSILON = enum.auto() # the model predicts epsilon - - -mean_type_dict = { - "previous_x": ModelMeanType.PREVIOUS_X, - "start_x": ModelMeanType.START_X, - "epsilon": ModelMeanType.EPSILON -} - - -class ModelVarType(enum.Enum): - """ - What is used as the model's output variance. - - The LEARNED_RANGE option has been added to allow the model to predict - values between FIXED_SMALL and FIXED_LARGE, making its job easier. - """ - - LEARNED = enum.auto() - FIXED_SMALL = enum.auto() - FIXED_LARGE = enum.auto() - LEARNED_RANGE = enum.auto() - - -var_type_dict = { - "learned": ModelVarType.LEARNED, - "fixed_small": ModelVarType.FIXED_SMALL, - "fixed_large": ModelVarType.FIXED_LARGE, - "learned_range": ModelVarType.LEARNED_RANGE -} - - -class GaussianDiffusion: - """ - Utilities for training and sampling diffusion modules. - - Ported directly from here, and then adapted over time to further experimentation. - https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 - - :param betas: a 1-D numpy array of betas for each diffusion timestep, - starting at T and going to 1. - :param model_mean_type: a ModelMeanType determining what the model outputs. - :param model_var_type: a ModelVarType determining how variance is output. - :param loss_type: a LossType determining the loss function to use. - :param rescale_timesteps: if True, pass floating point timesteps into the - model so that they are always scaled like in the - original paper (0 to 1000). - """ - - def __init__( - self, - *, - noise_schedule, - steps, - predict='start_x', - var_type="fixed_large", - losses="mse", # this can be a string or an array of strings - losses_multipliers=1., # this can be a single float or an array of floats - rescale_timesteps=False, - noise_std=1, - **kwargs - ): - assert predict in mean_type_dict.keys(), f"predict='{predict}' not supported" - self.model_mean_type = mean_type_dict[predict] - assert var_type in var_type_dict.keys(), f"var_type='{var_type}' not supported" - self.model_var_type = var_type_dict[var_type] - - # support for a linear combination (losses_multipliers)) of several loses - if isinstance(losses, str): - losses = [losses, ] # retro-compatibility - if isinstance(losses_multipliers, float): - losses_multipliers = [losses_multipliers, ] - assert len(losses) == len(losses_multipliers) - self.losses = losses - self.losses_multipliers = losses_multipliers - - self.rescale_timesteps = rescale_timesteps - self.noise_std = noise_std - - # Use float64 for accuracy. - betas = get_named_beta_schedule(noise_schedule, steps) - betas = np.array(betas, dtype=np.float64) - self.betas = betas - assert len(betas.shape) == 1, "betas must be 1-D" - assert (betas > 0).all() and (betas <= 1).all() - - self.num_timesteps = int(betas.shape[0]) - - alphas = 1.0 - betas - self.alphas_cumprod = np.cumprod(alphas, axis=0) - self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) - self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) - assert self.alphas_cumprod_prev.shape == (self.num_timesteps,) - - # calculations for diffusion q(x_t | x_{t-1}) and others - self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) - self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) - self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) - self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) # - self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) - - # calculations for posterior q(x_{t-1} | x_t, x_0) - self.posterior_variance = ( - betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) - ) # posterior_variance = beta_t * (1 - alpha_(t-1)^bar) / (1 - alpha_t^bar) - # log calculation clipped because the posterior variance is 0 at the - # beginning of the diffusion chain. - self.posterior_log_variance_clipped = np.log( - np.append(self.posterior_variance[1], self.posterior_variance[1:]) - ) - self.posterior_mean_coef1 = ( - betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) - ) - self.posterior_mean_coef2 = ( - (1.0 - self.alphas_cumprod_prev) - * np.sqrt(alphas) - / (1.0 - self.alphas_cumprod) - ) - - def q_mean_variance(self, x_start, t): - """ - Get the distribution q(x_t | x_0). - - :param x_start: the [N x C x ...] tensor of noiseless inputs. - :param t: the number of diffusion steps (minus 1). Here, 0 means one step. - :return: A tuple (mean, variance, log_variance), all of x_start's shape. - """ - mean = ( - _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start - ) - variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) - log_variance = _extract_into_tensor( - self.log_one_minus_alphas_cumprod, t, x_start.shape - ) - return mean, variance, log_variance - - def q_sample(self, x_start, t, noise=None): - """ - Diffuse the data for a given number of diffusion steps. - - In other words, sample from q(x_t | x_0). - -> using the reparametrization trick of: - sqrt(alfa) * x_0 + sqrt(1 - alfa) * eps - - :param x_start: the initial data batch. - :param t: the number of diffusion steps (minus 1). Here, 0 means one step. - :param noise: if specified, the split-out normal noise. - :return: A noisy version of x_start. - """ - if noise is None: - noise = th.randn_like(x_start) * self.noise_std - assert noise.shape == x_start.shape - weighed_x_start = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start - weighed_noise = _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise - return (weighed_x_start + weighed_noise) - - def q_posterior_mean_variance(self, x_start, x_t, t): - """ - Compute the mean and variance of the diffusion posterior: - - q(x_{t-1} | x_t, x_0) - - """ - assert x_start.shape == x_t.shape - posterior_mean = ( - _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start - + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t - ) - posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape) - posterior_log_variance_clipped = _extract_into_tensor( - self.posterior_log_variance_clipped, t, x_t.shape - ) - assert ( - posterior_mean.shape[0] - == posterior_variance.shape[0] - == posterior_log_variance_clipped.shape[0] - == x_start.shape[0] - ) - return posterior_mean, posterior_variance, posterior_log_variance_clipped - - def p_mean_variance( - self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None - ): - """ - Apply the model to get p(x_{t-1} | x_t), as well as a prediction of - the initial x, x_0. - - :param model: the model, which takes a signal and a batch of timesteps - as input. - :param x: the [N x C x ...] tensor at time t. - :param t: a 1-D Tensor of timesteps. - :param clip_denoised: if True, clip the denoised signal into [-1, 1]. - :param denoised_fn: if not None, a function which applies to the - x_start prediction before it is used to sample. Applies before - clip_denoised. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :return: a dict with the following keys: - - 'mean': the model mean output. - - 'variance': the model variance output. - - 'log_variance': the log of 'variance'. - - 'pred_xstart': the prediction for x_0. - """ - if model_kwargs is None: - model_kwargs = {} - - # for sampling (inference) - model_output = model.forward_with_cond_scale(x, self._scale_timesteps(t), model_kwargs) - # model_output = model(x, self._scale_timesteps(t), model_kwargs) - - if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]: - raise NotImplementedError("The ModelVarType {} is not implemented yet.".format(self.model_var_type)) - # assert model_output.shape == (B, C * 2, *x.shape[2:]) - # model_output, model_var_values = th.split(model_output, C, dim=1) - # if self.model_var_type == ModelVarType.LEARNED: - # model_log_variance = model_var_values - # model_variance = th.exp(model_log_variance) - # else: - # min_log = _extract_into_tensor( - # self.posterior_log_variance_clipped, t, x.shape - # ) - # max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) - # # The model_var_values is [-1, 1] for [min_var, max_var]. - # frac = (model_var_values + 1) / 2 - # model_log_variance = frac * max_log + (1 - frac) * min_log - # model_variance = th.exp(model_log_variance) - else: - model_variance, model_log_variance = { - # for fixedlarge, we set the initial (log-)variance like so - # to get a better decoder log likelihood. - ModelVarType.FIXED_LARGE: ( # default - np.append(self.posterior_variance[1], self.betas[1:]), - np.log(np.append(self.posterior_variance[1], self.betas[1:])), - ), - ModelVarType.FIXED_SMALL: ( - self.posterior_variance, - self.posterior_log_variance_clipped, - ), - }[self.model_var_type] - - model_variance = _extract_into_tensor(model_variance, t, x.shape) - model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) - - def process_xstart(x): - if denoised_fn is not None: - x = denoised_fn(x) - if clip_denoised: - return x.clamp(-1, 1) - return x - - # self.model_mean_type == ModelMeanType.START_X - if self.model_mean_type == ModelMeanType.PREVIOUS_X: - pred_xstart = process_xstart( - self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output) - ) - model_mean = model_output - # default the model predict START_X. - elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]: - if self.model_mean_type == ModelMeanType.START_X: - pred_xstart = process_xstart(model_output) - else: # model predict gaussian noise. - pred_xstart = process_xstart( - self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output) - ) - model_mean, _, _ = self.q_posterior_mean_variance( - x_start=pred_xstart, x_t=x, t=t - ) - else: - raise NotImplementedError(self.model_mean_type) - - assert ( - model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape - ) - return { - "mean": model_mean, - "variance": model_variance, - "log_variance": model_log_variance, - "pred_xstart": pred_xstart, - } - - def _predict_xstart_from_eps(self, x_t, t, eps): - assert x_t.shape == eps.shape - return ( - _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps - ) - - def _predict_xstart_from_xprev(self, x_t, t, xprev): - assert x_t.shape == xprev.shape - return ( # (xprev - coef2*x_t) / coef1 - _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev - - _extract_into_tensor( - self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape - ) - * x_t - ) - - def _predict_eps_from_xstart(self, x_t, t, pred_xstart): - return ( - _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - - pred_xstart - ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) - - def _scale_timesteps(self, t): - if self.rescale_timesteps: - return t.float() * (1000.0 / self.num_timesteps) - return t - - def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): - """ - Compute the mean for the previous step, given a function cond_fn that - computes the gradient of a conditional log probability with respect to - x. In particular, cond_fn computes grad(log(p(y|x))), and we want to - condition on y. - - This uses the conditioning strategy from Sohl-Dickstein et al. (2015). - """ - gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs) - new_mean = ( - p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() - ) - return new_mean - - def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): - """ - Compute what the p_mean_variance output would have been, should the - model's score function be conditioned by cond_fn. - - See condition_mean() for details on cond_fn. - - Unlike condition_mean(), this instead uses the conditioning strategy - from Song et al (2020). - """ - alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) - - eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"]) - eps = eps - (1 - alpha_bar).sqrt() * cond_fn( - x, self._scale_timesteps(t), **model_kwargs - ) - - out = p_mean_var.copy() - out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps) - out["mean"], _, _ = self.q_posterior_mean_variance( - x_start=out["pred_xstart"], x_t=x, t=t - ) - return out - - def p_sample( - self, - model, - x, - t, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - ): - """ - Sample x_{t-1} from the model at the given timestep. - - :param model: the model to sample from. - :param x: the current tensor at x_{t}. - :param t: the value of t, starting at 0 for the first diffusion step. - :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. - :param denoised_fn: if not None, a function which applies to the - x_start prediction before it is used to sample. - :param cond_fn: if not None, this is a gradient function that acts - similarly to the model. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :return: a dict containing the following keys: - - 'sample': a random sample from the model. - - 'pred_xstart': a prediction of x_0. - """ - out = self.p_mean_variance( - model, - x, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - model_kwargs=model_kwargs, - ) - noise = th.randn_like(x) * self.noise_std - nonzero_mask = ( - (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) - ) # no noise when t == 0 - if cond_fn is not None: - out["mean"] = self.condition_mean( - cond_fn, out, x, t, model_kwargs=model_kwargs - ) - noise_to_add = nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise - sample = out["mean"] + noise_to_add - return {"sample": sample, "pred_xstart": out["pred_xstart"], "noise_to_add": noise_to_add} - - def p_sample_loop( - self, - model, - shape, - noise=None, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - device=None, - progress=False, - max_step=-1, - ): - """ - Generate samples from the model. - - :param model: the model module. - :param shape: the shape of the samples, (N, C, H, W). - :param noise: if specified, the noise from the encoder to sample. - Should be of the same shape as `shape`. - :param clip_denoised: if True, clip x_start predictions to [-1, 1]. - :param denoised_fn: if not None, a function which applies to the - x_start prediction before it is used to sample. - :param cond_fn: if not None, this is a gradient function that acts - similarly to the model. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :param device: if specified, the device to create the samples on. - If not specified, use a model parameter's device. - :param progress: if True, show a tqdm progress bar. - :param max_step: last diffusion step that wants to perform. If -1, all steps are performed - :return: a non-differentiable batch of samples. - """ - final = None - for i, sample in enumerate(self.p_sample_loop_progressive( - model, - shape, - noise=noise, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - cond_fn=cond_fn, - model_kwargs=model_kwargs, - device=device, - progress=progress, - )): - if max_step != -1 and i == max_step: - break - final = sample - - return final - - def p_sample_loop_progressive( - self, - model, - shape, - noise=None, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - device=None, - progress=False, - pred=None - ): - """ - Generate samples from the model and yield intermediate samples from - each timestep of diffusion. - - Arguments are the same as p_sample_loop(). - Returns a generator over dicts, where each dict is the return value of - p_sample(). - """ - if device is None: - device = next(model.parameters()).device - assert isinstance(shape, (tuple, list)) - if noise is not None: - img = noise - else: - img = th.randn(*shape, device=device) * self.noise_std - indices = list(range(self.num_timesteps))[::-1] - - if progress: - # Lazy import so that we don't depend on tqdm. - from tqdm.auto import tqdm - - indices = tqdm(indices) - - for i in indices: - t = th.tensor([i] * shape[0], device=device) - with th.no_grad(): - out = self.p_sample( - model, - img, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - cond_fn=cond_fn, - model_kwargs=model_kwargs, - ) - yield out - img = out["sample"] - - def ddim_sample( - self, - model, - x, - t, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - eta=0.0, - ): - """ - Sample x_{t-1} from the model using DDIM. - - Same usage as p_sample(). - """ - out = self.p_mean_variance( - model, - x, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - model_kwargs=model_kwargs, - ) - - # out["mean"].shape == out["variance"].shape == out["log_variance"].shape == out["pred_xstart"].shape == (4800, 128) - # after broadcast, out["mean"].shape == out["variance"].shape - if cond_fn is not None: - out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) - - # Usually our model outputs epsilon, but we re-derive it (do the reverse of first term in Equation 12 in DDIM) - # in case we used x_start or x_prev prediction. - eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) # following Equation 12 in DDIM. - - alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) - alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) - sigma = ( - eta # hyperparameter - * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) - * th.sqrt(1 - alpha_bar / alpha_bar_prev) - ) - # following Equation 12 in DDIM. - noise = th.randn_like(x) * self.noise_std - mean_pred = ( - out["pred_xstart"] * th.sqrt(alpha_bar_prev) # x_pred * sqrt(alpha_(t-1)) - + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps - ) - nonzero_mask = ( - (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) - ) # no noise when t == 0 - sample = mean_pred + nonzero_mask * sigma * noise - # out["pred_xstart"] denotes 'predicted x_0' following Equation 12 in DDIM. - return {"sample": sample, "pred_xstart": out["pred_xstart"]} - - def ddim_reverse_sample( - self, - model, - x, - t, - clip_denoised=True, - denoised_fn=None, - model_kwargs=None, - eta=0.0, - ): - """ - Sample x_{t+1} from the model using DDIM reverse ODE. - """ - assert eta == 0.0, "Reverse ODE only for deterministic path" - out = self.p_mean_variance( - model, - x, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - model_kwargs=model_kwargs, - ) - # Usually our model outputs epsilon, but we re-derive it - # in case we used x_start or x_prev prediction. - eps = (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x - - out["pred_xstart"] - ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape) - alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape) - - # Equation 12. reversed - mean_pred = ( - out["pred_xstart"] * th.sqrt(alpha_bar_next) - + th.sqrt(1 - alpha_bar_next) * eps - ) - - return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]} - - def ddim_sample_loop( - self, - model, - shape, - noise=None, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - device=None, - progress=False, - eta=0.0, - ): - """ - Generate samples from the model using DDIM. - - Same usage as p_sample_loop(). - """ - final = None - for sample in self.ddim_sample_loop_progressive( - model, - shape, - noise=noise, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - cond_fn=cond_fn, - model_kwargs=model_kwargs, - device=device, - progress=progress, - eta=eta, - ): - final = sample - return final["sample"] - - def ddim_sample_loop_progressive( - self, - model, - shape, - noise=None, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - device=None, - progress=False, - eta=0.0, - ): - """ - Use DDIM to sample from the model and yield intermediate samples from - each timestep of DDIM. - - Same usage as p_sample_loop_progressive(). - """ - if device is None: - device = next(model.parameters()).device - assert isinstance(shape, (tuple, list)) - if noise is not None: - img = noise - else: - img = th.randn(*shape, device=device) * self.noise_std - indices = list(range(self.num_timesteps))[::-1] - - if progress: - # Lazy import so that we don't depend on tqdm. - from tqdm.auto import tqdm - - indices = tqdm(indices) - - for i in indices: - t = th.tensor([i] * shape[0], device=device) - with th.no_grad(): - out = self.ddim_sample( - model, - img, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - cond_fn=cond_fn, - model_kwargs=model_kwargs, - eta=eta, - ) - yield out - img = out["sample"] - - def denoise(self, model, x_start, t, model_kwargs=None, noise=None, all_kwargs=None): - raise NotImplementedError - - def get_gt( - self, - model, - obs, - pred - ): - raise NotImplementedError - - -def _extract_into_tensor(arr, timesteps, broadcast_shape): - """ - Extract values from a 1-D numpy array for a batch of indices. - - :param arr: the 1-D numpy array. - :param timesteps: a tensor of indices into the array to extract. - :param broadcast_shape: a larger shape of K dimensions with the batch - dimension equal to the length of timesteps. - :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. - """ - res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float() - while len(res.shape) < len(broadcast_shape): - res = res[..., None] - return res.expand(broadcast_shape) - - -class PriorLatentDiffusion(GaussianDiffusion): - def __init__(self, cfg, train_timesteps, inference_timesteps): - """ - :param cfg: diffusion_prior.scheduler - # conf.diffusion_prior.scheduler, - # conf.diffusion_prior.scheduler.num_train_timesteps, - # conf.diffusion_prior.scheduler.num_inference_timesteps, - """ - - super().__init__( - noise_schedule=cfg.get("noise_schedule", "cosine"), - steps=train_timesteps, - predict=cfg.get("predict", "start_x"), - var_type=cfg.get("var_type", "fixed_large"), - rescale_timesteps=cfg.get("rescale_timesteps", False), - noise_std=cfg.get("noise_std", 1) - ) - self.k = cfg.num_preds - - # for DDIM sampling - if cfg.timestep_spacing == "linspace": - self.indices = ( - np.linspace(0, train_timesteps - 1, inference_timesteps) - .round()[::-1] - .copy() - .astype(np.int64) - ) - elif cfg.timestep_spacing == "leading": - step_ratio = train_timesteps // inference_timesteps - self.indices = (np.arange(0, inference_timesteps) * step_ratio).round()[::-1].copy().astype(np.int64) - elif cfg.timestep_spacing == "trailing": - step_ratio = train_timesteps // inference_timesteps - self.indices = np.round(np.arange(train_timesteps, 0, -step_ratio)).astype(np.int64) - else: - self.indices = list(range(self.num_timesteps))[::-1] # [999, 998, ... 0] - - def ddim_sample( - self, - model, - x, - t, - t_prev, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - eta=0.0, - ): - """ - Sample x_{t-1} from the model using DDIM. - - Same usage as p_sample(). - """ - out = self.p_mean_variance( - model, - x, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - model_kwargs=model_kwargs, - ) - - # out["mean"].shape == out["variance"].shape == out["log_variance"].shape == out["pred_xstart"].shape == (4800, 128) - # after broadcast, out["mean"].shape == out["variance"].shape - if cond_fn is not None: - out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) - - # Usually our model outputs epsilon, but we re-derive it (do the reverse of first term in Equation 12 in DDIM) - # in case we used x_start or x_prev prediction. - eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) # following Equation 12 in DDIM. - - alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) - # alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) - alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod, t_prev, x.shape) - - sigma = ( - eta # hyperparameter - * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) - * th.sqrt(1 - alpha_bar / alpha_bar_prev) - ) - # following Equation 12 in DDIM. - noise = th.randn_like(x) * self.noise_std - mean_pred = ( - out["pred_xstart"] * th.sqrt(alpha_bar_prev) # x_pred * sqrt(alpha_(t-1)) - + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps - ) - nonzero_mask = ( - (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) - ) # no noise when t == 0 - sample = mean_pred + nonzero_mask * sigma * noise - # out["pred_xstart"] denotes 'predicted x_0' following Equation 12 in DDIM. - return {"sample": sample, "pred_xstart": out["pred_xstart"]} - - def ddim_sample_loop_progressive( - self, - matcher, - model, - noise=None, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - device=None, - progress=False, - eta=0.0, - shape=None, - ): - """ - Use DDIM to sample from the model and yield intermediate samples from - each timestep of DDIM. - - Same usage as p_sample_loop_progressive(). - """ - - # --------- - # copy of model_kwargs - model_kwargs = model_kwargs.copy() - - if device is None: - device = next(model.parameters()).device - assert isinstance(shape, (tuple, list)) - if noise is not None: - img = noise - else: - img = th.randn(*shape, device=device) * self.noise_std - - # indices = list(range(self.num_timesteps))[::-1] - indices = self.indices - - if progress: - # Lazy import so that we don't depend on tqdm. - from tqdm.auto import tqdm - indices = tqdm(indices) - - for idx, i in enumerate(indices): - t = th.tensor([i] * shape[0], device=device) # timestep inverse traversal - t_prev = th.tensor([indices[idx+1]] * shape[0], device=device) if idx < (len(indices)-1) else \ - th.tensor([0] * shape[0], device=device) - with th.no_grad(): - out = self.ddim_sample( - model, - img, - t, - t_prev, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - cond_fn=cond_fn, - model_kwargs=model_kwargs, - eta=eta, - ) - - out = { - "encoded_prediction": out["pred_xstart"], - "sample_enc": out["sample"], - } - - yield out - img = out["sample_enc"] - - def denoise(self, model, x_start, t, model_kwargs=None, noise=None, all_kwargs=None): - t = t.repeat_interleave(self.k, dim=0) # (2, 25) -> (2, ..., 2, 25, ..., 25) - bs = t.shape[0] - # t.shape: (batch_size * k, ) - - if model_kwargs is None: - model_kwargs = {} - if noise is None: - noise = th.randn_like(x_start) * self.noise_std # k different noises for each corresponding x_0 - - x_t = self.q_sample(x_start, t, noise=noise) # apply perturbations from '0' to 't' to original image (x_start) - model_output = model.forward(x_t, self._scale_timesteps(t), model_kwargs) - - target = { - ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(x_start=x_start, x_t=x_t, t=t)[0], - ModelMeanType.START_X: x_start, - ModelMeanType.EPSILON: noise, - }[self.model_mean_type] # here the target is set to x_0 (x_start) rather than noise (eps) - - results = { - "encoded_prediction": model_output, # encoded - "encoded_target": target, # encoded - } - # undo the repeat_interleave for the k predictions - results = {k: v.view(-1, self.k, *results[k].shape[1:]) for k, v in results.items()} - - return results - - -class DecoderLatentDiffusion(GaussianDiffusion): - def __init__(self, cfg, train_timesteps, inference_timesteps): - super().__init__( - noise_schedule=cfg.get("noise_schedule", "cosine"), - steps=train_timesteps, - predict=cfg.get("predict", "start_x"), - var_type=cfg.get("var_type", "fixed_large"), - rescale_timesteps=cfg.get("rescale_timesteps", False), - noise_std=cfg.get("noise_std", 1) - ) - self.k = cfg.num_preds - - # for DDIM sampling - if cfg.timestep_spacing == "linspace": - self.indices = ( - np.linspace(0, train_timesteps - 1, inference_timesteps) - .round()[::-1] - .copy() - .astype(np.int64) - ) - elif cfg.timestep_spacing == "leading": - step_ratio = train_timesteps // inference_timesteps - self.indices = (np.arange(0, inference_timesteps) * step_ratio).round()[::-1].copy().astype(np.int64) - elif cfg.timestep_spacing == "trailing": - step_ratio = train_timesteps // inference_timesteps - self.indices = np.round(np.arange(train_timesteps, 0, -step_ratio)).astype(np.int64) - else: - self.indices = list(range(self.num_timesteps))[::-1] # [999, 998, ... 0] - - def ddim_sample( - self, - model, - x, - t, - t_prev, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - eta=0.0, - ): - """ - Sample x_{t-1} from the model using DDIM. - - Same usage as p_sample(). - """ - out = self.p_mean_variance( - model, - x, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - model_kwargs=model_kwargs, - ) - - # out["mean"].shape == out["variance"].shape == out["log_variance"].shape == out["pred_xstart"].shape == (4800, 128) - # after broadcast, out["mean"].shape == out["variance"].shape - if cond_fn is not None: - out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) - - # Usually our model outputs epsilon, but we re-derive it (do the reverse of first term in Equation 12 in DDIM) - # in case we used x_start or x_prev prediction. - eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) # following Equation 12 in DDIM. - - alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) - # alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) - alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod, t_prev, x.shape) - - sigma = ( - eta # hyperparameter - * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) - * th.sqrt(1 - alpha_bar / alpha_bar_prev) - ) - # following Equation 12 in DDIM. - noise = th.randn_like(x) * self.noise_std - mean_pred = ( - out["pred_xstart"] * th.sqrt(alpha_bar_prev) # x_pred * sqrt(alpha_(t-1)) - + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps - ) - nonzero_mask = ( - (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) - ) # no noise when t == 0 - sample = mean_pred + nonzero_mask * sigma * noise - # out["pred_xstart"] denotes 'predicted x_0' following Equation 12 in DDIM. - return {"sample": sample, "pred_xstart": out["pred_xstart"]} - - def ddim_sample_loop_progressive( - self, - matcher, - model, - noise=None, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - device=None, - progress=False, - eta=0.0, - shape=None, - ): - """ - Use DDIM to sample from the model and yield intermediate samples from - each timestep of DDIM. - - Same usage as p_sample_loop_progressive(). - """ - self.shape = shape - # copy of model_kwargs - model_kwargs = model_kwargs.copy() - - if device is None: - device = next(model.parameters()).device - assert isinstance(shape, (tuple, list)) - if noise is not None: - img = noise - else: - img = th.randn(size=shape, device=device) * self.noise_std - - # indices = list(range(self.num_timesteps))[::-1] - indices = self.indices - - if progress: - # Lazy import so that we don't depend on tqdm. - from tqdm.auto import tqdm - indices = tqdm(indices) - - for idx, i in enumerate(indices): - t = th.tensor([i] * shape[0], device=device) # timestep inverse traversal - t_prev = th.tensor([indices[idx+1]] * shape[0], device=device) if idx < (len(indices)-1) else \ - th.tensor([0] * shape[0], device=device) - with th.no_grad(): - out = self.ddim_sample( - model, - img, - t, # t - t_prev, # t-1 - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - cond_fn=cond_fn, - model_kwargs=model_kwargs, - eta=eta, - ) - - out = { - "decoded_prediction": out["pred_xstart"], - "sample_enc": out["sample"], - } - - yield out - img = out["sample_enc"] - - def denoise(self, model, x_start, t, model_kwargs=None, noise=None, all_kwargs=None): - # t = t.repeat_interleave(self.k, dim=0) - - if model_kwargs is None: - model_kwargs = {} - if noise is None: - noise = th.randn_like(x_start) * self.noise_std # k different noises for each corresponding x_0 - - x_t = self.q_sample(x_start, t, noise=noise) # apply perturbations from '0' to 't' to original image (x_start) - model_output = model(x_t, self._scale_timesteps(t), model_kwargs) - - target = { - ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(x_start=x_start, x_t=x_t, t=t)[0], - ModelMeanType.START_X: x_start, - ModelMeanType.EPSILON: noise, - }[self.model_mean_type] # here the target is set to x_0 (x_start) or noise (epsilon) - - results = { - "prediction_emotion": model_output, - "target_emotion": target, - } - return results diff --git a/personalised/code/framework/motion_diffusion/diffusion/matchers.py b/personalised/code/framework/motion_diffusion/diffusion/matchers.py deleted file mode 100644 index c14cf2ddf21d53070e4ce0bc689a282c97d6dc58..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/matchers.py +++ /dev/null @@ -1,719 +0,0 @@ -""" -Code adapted from: -https://github.com/BarqueroGerman/BeLFusion -""" -from pathlib import Path -import hydra -import torch -import torch.nn as nn -import os -from einops import rearrange -from omegaconf import DictConfig -from hydra.utils import instantiate -from framework.motion_diffusion.diffusion.diffusion_decoder.transformer_denoiser import TransformerDenoiser, \ - lengths_to_mask -from framework.motion_diffusion.diffusion.diffusion_prior.transformer_prior import DiffusionPriorNetwork -from framework.motion_diffusion.diffusion.gaussian_diffusion import PriorLatentDiffusion, DecoderLatentDiffusion -from framework.motion_diffusion.diffusion.resample import UniformSampler -from framework.motion_diffusion.diffusion.rnn import LatentEmbedder -from framework.utils.util import from_pretrained_checkpoint, save_checkpoint - - -class EEGPredictionHead(nn.Module): - def __init__(self, input_dim, hidden_dim=256, output_dim=14, dropout=0.1): - super().__init__() - self.net = nn.Sequential( - nn.LayerNorm(input_dim), - nn.Linear(input_dim, hidden_dim), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(hidden_dim, output_dim), - ) - - def forward(self, x): - return self.net(x) - - def get_model_name(self): - return self.__class__.__name__ - - -class BaseLatentModel(nn.Module): - def __init__(self, cfg, emb_preprocessing=False, freeze_encoder=True, **kwargs): - super(BaseLatentModel, self).__init__() - self.emb_preprocessing = emb_preprocessing - self.freeze_encoder = freeze_encoder - def_dtype = torch.get_default_dtype() - - self.audio_encoder = instantiate(cfg.audio_encoder) - if cfg.latent_embedder is not None: - self.latent_embedder = instantiate(cfg.latent_embedder) - model_path = os.path.join(hydra.utils.get_original_cwd(), cfg.latent_embedder.checkpoint_path) - checkpoint = torch.load(model_path, map_location='cpu') - state_dict = checkpoint['state_dict'] - self.latent_embedder.load_state_dict(state_dict) - print(f"Successfully loaded latent embedder from {model_path}") - else: - self.latent_embedder = LatentEmbedder() - - if self.freeze_encoder: # freeze modules - for para in self.latent_embedder.parameters(): - para.requires_grad = False - - torch.set_default_dtype(def_dtype) - self.init_params = None - - def deepcopy(self): - assert self.init_params is not None, "Cannot deepcopy LatentUNetMatcher if init_params is None." - # I can't deep copy this class. I need to do this trick to make the deepcopy of everything - model_copy = self.__class__(**self.init_params) - weights_path = f'weights_temp_{id(model_copy)}.pt' - torch.save(self.state_dict(), weights_path) - model_copy.load_state_dict(torch.load(weights_path)) - os.remove(weights_path) - return model_copy - - def preprocess(self, emb): - stats = self.embed_emotion_stats - if stats is None: - return emb # when no checkpoint was loaded, there is no stats. - - if "standardize" in self.emb_preprocessing: - return (emb - stats["mean"]) / torch.sqrt(stats["var"]) - elif "normalize" in self.emb_preprocessing: - return 2 * (emb - stats["min"]) / (stats["max"] - stats["min"]) - 1 - elif "none" in self.emb_preprocessing.lower(): - return emb - else: - raise NotImplementedError(f"Error on the embedding preprocessing value: '{self.emb_preprocessing}'") - - def undo_preprocess(self, emb): - stats = self.embed_emotion_stats - if stats is None: - return emb # when no checkpoint was loaded, there is no stats. - - if "standardize" in self.emb_preprocessing: - return torch.sqrt(stats["var"]) * emb + stats["mean"] - elif "normalize" in self.emb_preprocessing: - return (emb + 1) * (stats["max"] - stats["min"]) / 2 + stats["min"] - elif "none" in self.emb_preprocessing.lower(): - return emb - else: - raise NotImplementedError(f"Error on the embedding preprocessing value: '{self.emb_preprocessing}'") - - def forward(self, pred, timesteps, seq_em): - raise NotImplementedError("This is an abstract class.") - - # override checkpointing - def state_dict(self): - return self.model.state_dict() - - def load_state_dict(self, state_dict): - self.model.load_state_dict(state_dict) - - def to(self, device): - self.model = self.model.to(device) - return self - - def cuda(self): - return self.to(torch.device("cuda")) - - # override eval and train - def train(self, mode=True): - self.model.train(mode) - - def eval(self): - self.model.eval() - - -class PriorLatentMatcher(BaseLatentModel): - def __init__(self, - conf: DictConfig = None, - module_dict_cfg: DictConfig = None, - stage: str = 'fit', - task: str = 'online', - **kwargs): - cfg = conf.args - super(PriorLatentMatcher, self).__init__( - module_dict_cfg, - emb_preprocessing=cfg.emb_preprocessing, - freeze_encoder=cfg.freeze_encoder, - **kwargs, - ) - - self.stage = stage - self.task = task - self.token_len = cfg.token_len - self.window_size = cfg.get("window_size", 30) - self.s_ratio = cfg.get("s_ratio", 2) - self.s_window_size = cfg.get("s_window_size", self.window_size * self.s_ratio) - - self.init_params = { - "audio_dim": cfg.get("audio_dim", 768), - "window_size": self.s_window_size, - "_3dmm_dim": cfg.get("_3dmm_dim", 58), - "speaker_emb_dim": cfg.get("speaker_emb_dim", 512), - "latent_dim": cfg.get("latent_dim", 512), - "depth": cfg.get("depth", 4), - "num_time_layers": cfg.get("num_time_layers", 2), - "num_time_embeds": cfg.get("num_time_embeds", 1), - "num_time_emb_channels": cfg.get("num_time_emb_channels", 64), - "time_last_act": cfg.get("time_last_act", False), - "use_learned_query": cfg.get("use_learned_query", True), - "s_audio_cond_drop_prob": cfg.get("s_audio_cond_drop_prob", 0.2), - "s_latentemb_cond_drop_prob": cfg.get("s_latentemb_cond_drop_prob", 1.0), - "s_3dmm_cond_drop_prob": cfg.get("s_3dmm_cond_drop_prob", 0.2), - "guidance_scale": cfg.get("guidance_scale", 1.0), - "dim_head": cfg.get("dim_head", 64), - "heads": cfg.get("heads", 8), - "ff_mult": cfg.get("ff_mult", 4), - "norm_in": cfg.get("norm_in", False), - "norm_out": cfg.get("norm_out", True), - "attn_dropout": cfg.get("attn_dropout", 0.0), - "ff_dropout": cfg.get("ff_dropout", 0.0), - "final_proj": cfg.get("final_proj", True), - "normformer": cfg.get("normformer", False), - "rotary_emb": cfg.get("rotary_emb", True), - } - self.model = DiffusionPriorNetwork(**self.init_params) - - self.prior_diffusion = PriorLatentDiffusion( - conf.scheduler, - conf.scheduler.num_train_timesteps, - conf.scheduler.num_inference_timesteps, - ) - self.schedule_sampler = UniformSampler(self.prior_diffusion) - self.num_preds = conf.scheduler.num_preds - - def _select_training_windows(self, - speaker_audio_input, - speaker_emotion_input, - speaker_3dmm_input, - listener_emotion_input): - speaker_len = speaker_audio_input.shape[1] - listener_len = listener_emotion_input.shape[1] - - if speaker_len > self.s_window_size: - max_start = speaker_len - self.s_window_size - window_start = torch.randint(0, max_start + 1, (1,), device=speaker_audio_input.device).item() - else: - window_start = 0 - - speaker_audio_input = speaker_audio_input[:, window_start:window_start + self.s_window_size] - speaker_emotion_input = speaker_emotion_input[:, window_start:window_start + self.s_window_size] - speaker_3dmm_input = speaker_3dmm_input[:, window_start:window_start + self.s_window_size] - - listener_start = min(window_start, max(listener_len - self.window_size, 0)) - listener_emotion_input = listener_emotion_input[:, listener_start:listener_start + self.window_size] - - return speaker_audio_input, speaker_emotion_input, speaker_3dmm_input, listener_emotion_input - - def _forward( - self, - speaker_audio_input=None, - speaker_emotion_input=None, - speaker_3dmm_input=None, - listener_emotion_input=None, - **kwargs, - ): - if self.stage == "test": - raise RuntimeError("PriorLatentMatcher is trained and checkpointed only; inference still uses decoder output.") - - (speaker_audio_input, - speaker_emotion_input, - speaker_3dmm_input, - listener_emotion_input) = self._select_training_windows( - speaker_audio_input, - speaker_emotion_input, - speaker_3dmm_input, - listener_emotion_input, - ) - - with torch.no_grad(): - s_audio_encodings = self.audio_encoder._encode(speaker_audio_input) - s_audio_encodings = s_audio_encodings.repeat_interleave(self.num_preds, dim=0) - - s_latent_embed = self.latent_embedder.encode(speaker_emotion_input).unsqueeze(1) - s_latent_embed = s_latent_embed.repeat_interleave(self.num_preds, dim=0) - - s_3dmm_encodings = speaker_3dmm_input.repeat_interleave(self.num_preds, dim=0) - - listener_latent_embed = self.latent_embedder.encode(listener_emotion_input).unsqueeze(1) - listener_latent_embed = listener_latent_embed.repeat_interleave(self.num_preds, dim=0) - - model_kwargs = { - "speaker_audio_encodings": s_audio_encodings, - "speaker_latent_emb": s_latent_embed, - "speaker_3dmm_encodings": s_3dmm_encodings, - } - - t, _ = self.schedule_sampler.sample(listener_latent_embed.shape[0] // self.num_preds, - listener_latent_embed.device) - output_prior = self.prior_diffusion.denoise( - self.model, - listener_latent_embed, - t, - model_kwargs=model_kwargs, - ) - return output_prior - - def forward(self, **kwargs): - return self._forward(**kwargs) - - -class DecoderLatentMatcher(BaseLatentModel): - def __init__(self, - conf: DictConfig = None, - module_dict_cfg: DictConfig = None, - stage: str = 'fit', - task: str = 'online', - **kwargs): - cfg = conf.args - super(DecoderLatentMatcher, self).__init__( - module_dict_cfg, - emb_preprocessing=cfg.emb_preprocessing, - freeze_encoder=cfg.freeze_encoder, - **kwargs, - ) - - self.stage = stage - self.task = task - self.token_len = cfg.token_len - self.window_size = cfg.get("window_size", 30) - self.s_ratio = cfg.get("s_ratio", 2) - self.emotion_dim = cfg.get("nfeats", 25) - self.encode_emotion = cfg.get("encode_emotion", False) - self.encode_3dmm = cfg.get("encode_3dmm", False) - - self.init_params = { - "task": task, - "window_size": self.window_size, - "encode_emotion": self.encode_emotion, - "encode_3dmm": self.encode_3dmm, - "ablation_skip_connection": cfg.get("ablation_skip_connection", True), - "nfeats": cfg.get("nfeats", 25), - "latent_dim": cfg.get("latent_dim", 512), - "ff_size": cfg.get("ff_size", 1024), - "num_layers": cfg.get("num_layers", 6), - "num_heads": cfg.get("num_heads", 4), - "dropout": cfg.get("dropout", 0.1), - "normalize_before": cfg.get("normalize_before", False), - "activation": cfg.get("activation", "gelu"), - "flip_sin_to_cos": cfg.get("flip_sin_to_cos", True), - "return_intermediate_dec": cfg.get("return_intermediate_dec", False), - "position_embedding": cfg.get("position_embedding", "learned"), - "arch": cfg.get("arch", "trans_enc"), - "freq_shift": cfg.get("freq_shift", 0), - "time_encoded_dim": cfg.get("time_encoded_dim", 64), - "s_audio_dim": cfg.get("s_audio_dim", 768), - "s_audio_scale": cfg.get("s_audio_scale", cfg.get("latent_dim", 512) ** -0.5), - "s_emotion_dim": cfg.get("s_emotion_dim", 25), - "s_3dmm_dim": cfg.get("s_3dmm_dim", 58), - "concat": cfg.get("concat", "concat_first"), - "condition_concat": cfg.get("condition_concat", "token_concat"), - "guidance_scale": cfg.get("guidance_scale", 7.5), - "s_audio_enc_drop_prob": cfg.get("s_audio_enc_drop_prob", 0.2), - "s_latent_embed_drop_prob": cfg.get("s_latent_embed_drop_prob", 0.2), - "s_3dmm_enc_drop_prob": cfg.get("s_3dmm_enc_drop_prob", 0.2), - "s_emotion_enc_drop_prob": cfg.get("s_emotion_enc_drop_prob", 1.0), - "past_l_emotion_drop_prob": cfg.get("past_l_emotion_drop_prob", 1.0), - } - self.use_past_frames = cfg.get("use_past_frames", False) - - self.model = TransformerDenoiser(**self.init_params) - - self.decoder_diffusion = DecoderLatentDiffusion( - conf.scheduler, - conf.scheduler.num_train_timesteps, - conf.scheduler.num_inference_timesteps, - ) - self.schedule_sampler = UniformSampler(self.decoder_diffusion) - self.num_preds = conf.scheduler.num_preds - - def _forward( - self, - speaker_audio_input=None, - speaker_emotion_input=None, - speaker_3dmm_input=None, - listener_emotion_input=None, - past_listener_emotion=None, - motion_length=None, - ): - with torch.no_grad(): - s_audio_encodings = self.audio_encoder._encode(speaker_audio_input) - s_audio_encodings = s_audio_encodings.repeat_interleave(self.num_preds, dim=0) - - # freeze latent RNN_VAE embedder to extract speaker latent embedding - s_latent_embed = self.latent_embedder.encode(speaker_emotion_input).unsqueeze(1) - s_latent_embed = s_latent_embed.repeat_interleave(self.num_preds, dim=0) - # shape: (batch_size * num_preds, 1, ...) - - # s_3dmm_encodings = self.latent_3dmm_embedder.get_encodings(speaker_3dmm_input) - s_3dmm_encodings = speaker_3dmm_input.repeat_interleave(self.num_preds, dim=0) - # shape: (bs * num_preds, s_w, ...) - - s_emotion_encodings = speaker_emotion_input.repeat_interleave(self.num_preds, dim=0) - # shape: (bs * num_preds, s_w, ...) - - # past arrives either per-sample (bs, l_w, d) [training / parallel test] or - # already expanded to (bs*num_preds, l_w, d) [autoregressive online test, where - # each prediction continues its OWN previous window]; only expand the former. - if past_listener_emotion is not None and \ - past_listener_emotion.shape[0] != s_audio_encodings.shape[0]: - past_listener_emotion = past_listener_emotion.repeat_interleave( - self.num_preds, dim=0) - # shape: (bs * num_preds, l_w, ...) - - motion_length = motion_length.repeat_interleave( - self.num_preds, dim=0) if motion_length is not None else None - - model_kwargs = { - "speaker_audio_encodings": s_audio_encodings, - "speaker_latent_embed": s_latent_embed, - "speaker_3dmm_encodings": s_3dmm_encodings, - "speaker_emotion_encodings": s_emotion_encodings, - "past_listener_emotion": past_listener_emotion, - "motion_length": motion_length, - } - - if self.stage == "test": - bs, l, _ = s_audio_encodings.shape # bz * num_preds - with torch.no_grad(): - output = [output for output in self.decoder_diffusion.ddim_sample_loop_progressive( - matcher=self, - model=self.model, - model_kwargs=model_kwargs, - shape=(bs, self.window_size if self.task == "online" else l, self.emotion_dim), - )][-1] # get last output - - output_listener_emotion = output["sample_enc"] # (bz * num_preds, l_w, d=25) - output_listener_emotion = rearrange(output_listener_emotion, - "(b n) w d -> b n w d", n=self.num_preds) - output_whole = {"prediction_emotion": output_listener_emotion} - - else: - listener_emotion_input = listener_emotion_input.repeat_interleave(self.num_preds, dim=0) - x_start_selected = listener_emotion_input # (bs * num_preds, l_w, ...) - - t, _ = self.schedule_sampler.sample(x_start_selected.shape[0], x_start_selected.device) - timesteps = t.long() - - output_whole = self.decoder_diffusion.denoise(self.model, x_start_selected, timesteps, - model_kwargs=model_kwargs) - if motion_length is not None: # offline task zero masking - device = x_start_selected.get_device() - output_mask = lengths_to_mask(motion_length, device=device, max_len=x_start_selected.shape[1]) - # print(f'output_whole["prediction_emotion"] shape: {output_whole["prediction_emotion"].shape}') - output_whole["prediction_emotion"] = (output_whole["prediction_emotion"] - * output_mask.float().unsqueeze(-1)) - - output_whole = {k: v.view(-1, self.num_preds, *output_whole[k].shape[1:]) for k, v in output_whole.items()} - return output_whole - - def forward(self, **kwargs): - return self._forward(**kwargs) - - -class LatentMatcher(nn.Module): - def __init__(self, - task: str = "online", - stage: str = "fit", - device: str = None, - diffusion_prior: DictConfig = None, - diffusion_decoder: DictConfig = None, - latent_embedder: DictConfig = None, - audio_encoder: DictConfig = None, - eeg_head: DictConfig = None, - resumed_training: bool = False, - auto_load_ckpt: bool = True, - **kwargs): - super().__init__() - if device is None: - device = "cuda" if torch.cuda.is_available() else "cpu" - self.task = task - self.stage = stage - self.kwargs = kwargs - - module_dict_cfg = DictConfig( - {"latent_embedder": latent_embedder, - "audio_encoder": audio_encoder,} - ) - - self.diffusion_prior_cfg = diffusion_prior - self.diffusion_prior = None - if self.diffusion_prior_cfg is not None: - self.diffusion_prior = PriorLatentMatcher(self.diffusion_prior_cfg, - task=task, - stage=stage, - module_dict_cfg=module_dict_cfg, - **kwargs) - - self.diffusion_decoder_cfg = diffusion_decoder - self.diffusion_decoder = DecoderLatentMatcher(self.diffusion_decoder_cfg, - task=task, - stage=stage, - module_dict_cfg=module_dict_cfg, - **kwargs) - self.eeg_head = None - self.eeg_head_pooling = "mean" - self.eeg_detach_prediction_emotion = True - self.eeg_use_speaker_audio = True - self.eeg_use_speaker_emotion = True - self.eeg_use_speaker_3dmm = True - self.eeg_use_prediction_emotion = True - self.eeg_speaker_audio_dim = 0 - self.eeg_speaker_emotion_dim = 0 - self.eeg_speaker_3dmm_dim = 0 - self.eeg_prediction_emotion_dim = 0 - if eeg_head is not None and eeg_head.get("enabled", False): - decoder_args = self.diffusion_decoder_cfg.args - self.eeg_head_pooling = eeg_head.get("pooling", "mean") - self.eeg_detach_prediction_emotion = eeg_head.get("detach_prediction_emotion", True) - self.eeg_use_speaker_audio = eeg_head.get("use_speaker_audio", True) - self.eeg_use_speaker_emotion = eeg_head.get("use_speaker_emotion", True) - self.eeg_use_speaker_3dmm = eeg_head.get("use_speaker_3dmm", True) - self.eeg_use_prediction_emotion = eeg_head.get("use_prediction_emotion", True) - self.eeg_speaker_audio_dim = decoder_args.get("s_audio_dim", 768) \ - if self.eeg_use_speaker_audio else 0 - self.eeg_speaker_emotion_dim = decoder_args.get("s_emotion_dim", 25) \ - if self.eeg_use_speaker_emotion else 0 - self.eeg_speaker_3dmm_dim = decoder_args.get("s_3dmm_dim", 58) \ - if self.eeg_use_speaker_3dmm else 0 - self.eeg_prediction_emotion_dim = decoder_args.get("nfeats", 25) \ - if self.eeg_use_prediction_emotion else 0 - eeg_input_dim = ( - self.eeg_speaker_audio_dim - + self.eeg_speaker_emotion_dim - + self.eeg_speaker_3dmm_dim - + self.eeg_prediction_emotion_dim - ) - if eeg_input_dim <= 0: - raise ValueError("At least one EEG head input source must be enabled.") - self.eeg_head = EEGPredictionHead( - input_dim=eeg_head.get("input_dim", eeg_input_dim), - hidden_dim=eeg_head.get("hidden_dim", 256), - output_dim=eeg_head.get("output_dim", 14), - dropout=eeg_head.get("dropout", 0.5), - ) - load_ckpt = False - want_last = False - want_best = False - - if resumed_training: - load_ckpt = True - want_last = True - if stage == "test": - load_ckpt = True - want_best = True - - if load_ckpt and auto_load_ckpt: - ckpt_path = self.get_ckpt_path( - self.diffusion_decoder.model, - runid="resume_runid", - epoch=None, - best=want_best, - last=want_last, - ) - from_pretrained_checkpoint(str(ckpt_path), self.diffusion_decoder.model, device) - if self.diffusion_prior is not None: - prior_ckpt_path = self.get_ckpt_path( - self.diffusion_prior.model, - runid="resume_runid", - epoch=None, - best=want_best, - last=want_last, - create_dir=False, - ) - if os.path.exists(prior_ckpt_path): - from_pretrained_checkpoint(str(prior_ckpt_path), self.diffusion_prior.model, device) - elif resumed_training: - raise FileNotFoundError(f"Missing prior checkpoint for resumed training: {prior_ckpt_path}") - if self.eeg_head is not None: - eeg_ckpt_path = self.get_ckpt_path( - self.eeg_head, - runid="resume_runid", - epoch=None, - best=want_best, - last=want_last, - create_dir=False, - ) - if os.path.exists(eeg_ckpt_path): - from_pretrained_checkpoint(str(eeg_ckpt_path), self.eeg_head, device) - elif resumed_training: - raise FileNotFoundError(f"Missing EEG head checkpoint for resumed training: {eeg_ckpt_path}") - - def freeze_except_eeg_head(self): - if self.eeg_head is None: - raise RuntimeError("Cannot train EEG head only because eeg_head is disabled.") - - for parameter in self.parameters(): - parameter.requires_grad = False - for parameter in self.eeg_head.parameters(): - parameter.requires_grad = True - - def set_eeg_head_train_mode(self): - if self.eeg_head is None: - raise RuntimeError("Cannot train EEG head only because eeg_head is disabled.") - - self.eval() - self.eeg_head.train() - - def _pool_eeg_sequence(self, feature, expected_dim, batch_size, num_preds, device, dtype): - if expected_dim <= 0: - return None - if feature is None or feature.numel() == 0: - return torch.zeros(batch_size, num_preds, expected_dim, device=device, dtype=dtype) - - feature = feature.to(device=device, dtype=dtype) - if feature.dim() == 3: - if self.eeg_head_pooling == "last": - pooled = feature[:, -1] - elif self.eeg_head_pooling == "mean": - pooled = feature.mean(dim=1) - else: - raise ValueError(f"Unknown EEG head pooling: {self.eeg_head_pooling}") - elif feature.dim() == 2: - pooled = feature - else: - raise ValueError(f"Unsupported EEG condition shape: {feature.shape}") - - return pooled.unsqueeze(1).expand(-1, num_preds, -1) - - def _attach_eeg_outputs(self, outputs, speaker_audio_input=None, speaker_emotion_input=None, - speaker_3dmm_input=None, listener_eeg_input=None, listener_eeg_mask=None): - if self.eeg_head is None: - return outputs - - prediction_emotion = outputs.get("prediction_emotion") - if prediction_emotion is None: - return outputs - - batch_size, num_preds = prediction_emotion.shape[:2] - device = prediction_emotion.device - dtype = prediction_emotion.dtype - feature_list = [] - - speaker_audio_feature = self._pool_eeg_sequence( - speaker_audio_input, self.eeg_speaker_audio_dim, batch_size, num_preds, device, dtype) - if speaker_audio_feature is not None: - feature_list.append(speaker_audio_feature) - - speaker_emotion_feature = self._pool_eeg_sequence( - speaker_emotion_input, self.eeg_speaker_emotion_dim, batch_size, num_preds, device, dtype) - if speaker_emotion_feature is not None: - feature_list.append(speaker_emotion_feature) - - speaker_3dmm_feature = self._pool_eeg_sequence( - speaker_3dmm_input, self.eeg_speaker_3dmm_dim, batch_size, num_preds, device, dtype) - if speaker_3dmm_feature is not None: - feature_list.append(speaker_3dmm_feature) - - if self.eeg_head_pooling == "last": - pooled_emotion = prediction_emotion[:, :, -1] - elif self.eeg_head_pooling == "mean": - pooled_emotion = prediction_emotion.mean(dim=2) - else: - raise ValueError(f"Unknown EEG head pooling: {self.eeg_head_pooling}") - if self.eeg_detach_prediction_emotion: - pooled_emotion = pooled_emotion.detach() - if self.eeg_use_prediction_emotion: - feature_list.append(pooled_emotion) - - prediction_eeg = self.eeg_head(torch.cat(feature_list, dim=-1)) - outputs["prediction_eeg"] = prediction_eeg - - if listener_eeg_input is None: - return outputs - - target_eeg = listener_eeg_input.to(prediction_eeg.device).float() - target_eeg_mask = listener_eeg_mask.to(prediction_eeg.device).float() \ - if listener_eeg_mask is not None else torch.ones_like(target_eeg) - if target_eeg.dim() == 2: - target_eeg = target_eeg.unsqueeze(1).expand(-1, num_preds, -1) - if target_eeg_mask.dim() == 2: - target_eeg_mask = target_eeg_mask.unsqueeze(1).expand(-1, num_preds, -1) - outputs["target_eeg"] = target_eeg - outputs["target_eeg_mask"] = target_eeg_mask - return outputs - - def forward( - self, - speaker_audio_input=None, - speaker_emotion_input=None, - speaker_3dmm_input=None, - listener_emotion_input=None, - listener_eeg_input=None, - listener_eeg_mask=None, - past_listener_emotion=None, - motion_length=None, - ): - - outputs = self.diffusion_decoder.forward( - speaker_audio_input=speaker_audio_input, - speaker_emotion_input=speaker_emotion_input, - speaker_3dmm_input=speaker_3dmm_input, - listener_emotion_input=listener_emotion_input, - past_listener_emotion=past_listener_emotion, - motion_length=motion_length, - ) - outputs = self._attach_eeg_outputs( - outputs, - speaker_audio_input=speaker_audio_input, - speaker_emotion_input=speaker_emotion_input, - speaker_3dmm_input=speaker_3dmm_input, - listener_eeg_input=listener_eeg_input, - listener_eeg_mask=listener_eeg_mask, - ) - # outputs['prediction_emotion']: (bz, num_preds, s_w, emotion_dim) - if self.stage == "test" or self.diffusion_prior is None: - return outputs - - output_prior = self.diffusion_prior.forward( - speaker_audio_input=speaker_audio_input, - speaker_emotion_input=speaker_emotion_input, - speaker_3dmm_input=speaker_3dmm_input, - listener_emotion_input=listener_emotion_input, - ) - - return { - "output_prior": output_prior, - "output_decoder": outputs, - } - - def get_ckpt_path(self, model, runid="current_runid", epoch=None, best=False, last=False, create_dir=True): - ckpt_dir = Path(hydra.utils.to_absolute_path(self.kwargs.get("ckpt_dir"))) - run_id = Path(self.kwargs.get(runid)) - ckpt_dir = str(ckpt_dir / run_id / model.get_model_name()) - if create_dir: - os.makedirs(ckpt_dir, exist_ok=True) - - ckpt_path = None - if epoch is not None: - ckpt_path = os.path.join(ckpt_dir, f"checkpoint_{epoch}.pth") - if best: - ckpt_path = os.path.join(ckpt_dir, "checkpoint_best.pth") - if last: - ckpt_path = os.path.join(ckpt_dir, "checkpoint_last.pth") - assert ckpt_path is not None, "No checkpoint path is provided." - return ckpt_path - - def save_ckpt(self, optimizer, epoch=None, best=False, last=False, best_loss=float("inf")): - models = [self.diffusion_decoder.model] - if self.diffusion_prior is not None: - models.append(self.diffusion_prior.model) - if self.eeg_head is not None: - models.append(self.eeg_head) - - for model in models: - ckpt_path = self.get_ckpt_path(model, epoch=epoch, best=best, last=last) - save_checkpoint(ckpt_path, model, optimizer, epoch=epoch, best_loss=best_loss) - - def obtain_shapes(self, modified_layers): - shape_dict = {} - for name, module in self.named_modules(): - if name not in modified_layers: - continue - if hasattr(module, "weight"): - shape_dict[name] = torch.tensor(module.weight.size()) - elif hasattr(module, "in_proj_weight"): - shape_dict[name] = torch.tensor(module.in_proj_weight.size()) - return shape_dict diff --git a/personalised/code/framework/motion_diffusion/diffusion/matchers_causal.py b/personalised/code/framework/motion_diffusion/diffusion/matchers_causal.py deleted file mode 100644 index 8219b5c990324e9d256e904a8ec8b2e51015b473..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/matchers_causal.py +++ /dev/null @@ -1,125 +0,0 @@ -""" -Causal / coarse-to-fine matcher wrappers -========================================= -Non-invasive subclasses that swap the baseline `TransformerDenoiser` for the -`CausalTransformerDenoiser` and surface the coarse 8-class logits to the loss. - - * `CoarseDecoderLatentMatcher` rebuilds `self.model` as the causal denoiser - (reusing the already-resolved `init_params`) and, during training, copies - the stashed `_coarse_logits` into the output dict. - * `CausalLatentMatcher` is a thin `LatentMatcher` that uses the coarse decoder - wrapper while keeping the prior / EEG head / checkpoint loading identical. - -No original file is modified; select these via Hydra `_target_` in a new config. -""" -import os - -from omegaconf import DictConfig - -from framework.motion_diffusion.diffusion.matchers import ( - DecoderLatentMatcher, - LatentMatcher, -) -from framework.motion_diffusion.diffusion.diffusion_decoder.transformer_denoiser_causal import ( - CausalTransformerDenoiser, -) -from framework.utils.util import from_pretrained_checkpoint - - -class CoarseDecoderLatentMatcher(DecoderLatentMatcher): - def __init__(self, conf: DictConfig = None, **kwargs): - super().__init__(conf, **kwargs) - cfg = conf.args - coarse_kwargs = dict( - lag_max=int(cfg.get("lag_max", 60)), - lag_lookahead=int(cfg.get("lag_lookahead", 0)), - coarse_classes=int(cfg.get("coarse_classes", 8)), - coarse_hidden=int(cfg.get("coarse_hidden", 256)), - coarse_emo_start=int(cfg.get("coarse_emo_start", 17)), - use_lag_bias=bool(cfg.get("use_lag_bias", True)), - use_coarse=bool(cfg.get("use_coarse", True)), - ) - # Rebuild the denoiser as the causal+coarse variant, reusing the - # init_params resolved by the parent (latent_dim, num_heads, drop probs ...). - self.model = CausalTransformerDenoiser(**self.init_params, **coarse_kwargs) - - def _forward(self, **kwargs): - out = super()._forward(**kwargs) - if self.stage != "test": - coarse_logits = getattr(self.model, "_coarse_logits", None) - if coarse_logits is not None: - # (bs*num_preds, T, C) -> (bs, num_preds, T, C) - out["coarse_logits"] = coarse_logits.view( - -1, self.num_preds, *coarse_logits.shape[1:]) - return out - - -class CausalLatentMatcher(LatentMatcher): - def __init__( - self, - task: str = "online", - stage: str = "fit", - device: str = None, - diffusion_prior: DictConfig = None, - diffusion_decoder: DictConfig = None, - latent_embedder: DictConfig = None, - audio_encoder: DictConfig = None, - eeg_head: DictConfig = None, - resumed_training: bool = False, - auto_load_ckpt: bool = True, - **kwargs, - ): - import torch - - # Build everything via the parent but defer checkpoint loading so we can - # swap in the coarse decoder first. - super().__init__( - task=task, stage=stage, device=device, - diffusion_prior=diffusion_prior, diffusion_decoder=diffusion_decoder, - latent_embedder=latent_embedder, audio_encoder=audio_encoder, - eeg_head=eeg_head, resumed_training=resumed_training, - auto_load_ckpt=False, **kwargs, - ) - if device is None: - device = "cuda" if torch.cuda.is_available() else "cpu" - - module_dict_cfg = DictConfig( - {"latent_embedder": latent_embedder, "audio_encoder": audio_encoder}) - self.diffusion_decoder = CoarseDecoderLatentMatcher( - self.diffusion_decoder_cfg, task=task, stage=stage, - module_dict_cfg=module_dict_cfg, **kwargs) - - if auto_load_ckpt: - self._load_causal_checkpoints(device, resumed_training, stage) - - def _load_causal_checkpoints(self, device, resumed_training, stage): - load_ckpt = bool(resumed_training) or stage == "test" - if not load_ckpt: - return - want_last = bool(resumed_training) - want_best = stage == "test" - - ckpt_path = self.get_ckpt_path( - self.diffusion_decoder.model, runid="resume_runid", - epoch=None, best=want_best, last=want_last) - from_pretrained_checkpoint(str(ckpt_path), self.diffusion_decoder.model, device) - - if self.diffusion_prior is not None: - prior_ckpt_path = self.get_ckpt_path( - self.diffusion_prior.model, runid="resume_runid", - epoch=None, best=want_best, last=want_last, create_dir=False) - if os.path.exists(prior_ckpt_path): - from_pretrained_checkpoint(str(prior_ckpt_path), self.diffusion_prior.model, device) - elif resumed_training: - raise FileNotFoundError( - f"Missing prior checkpoint for resumed training: {prior_ckpt_path}") - - if self.eeg_head is not None: - eeg_ckpt_path = self.get_ckpt_path( - self.eeg_head, runid="resume_runid", - epoch=None, best=want_best, last=want_last, create_dir=False) - if os.path.exists(eeg_ckpt_path): - from_pretrained_checkpoint(str(eeg_ckpt_path), self.eeg_head, device) - elif resumed_training: - raise FileNotFoundError( - f"Missing EEG head checkpoint for resumed training: {eeg_ckpt_path}") diff --git a/personalised/code/framework/motion_diffusion/diffusion/matchers_velocity.py b/personalised/code/framework/motion_diffusion/diffusion/matchers_velocity.py deleted file mode 100644 index f0bd206541468bcd048e1c68de305d449f8e8675..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/matchers_velocity.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Velocity-space variant of the PerFRDiff latent matcher. - -Only the *decoder* diffusion is reparametrised to operate in frame-delta -(velocity) space; the diffusion prior, EEG head and everything else are left -exactly as in the baseline ``matchers.py`` (the original file is untouched). - -Rationale: the listener reaction is modelled by predicting frame-to-frame change -(velocity) rather than absolute emotion values. In CCC terms (the FRC metric) this -directly pushes the temporal-correlation term and discourages variance collapse, -and it tends to yield smoother, more natural reactions — analogous to predicting -epsilon/v instead of x0 in standard diffusion. - -The only behavioural change lives in ``VelocityDecoderLatentMatcher._forward``: - * training: the diffusion target ``x_start`` becomes ``to_delta(listener_emotion)`` - so the network is supervised on velocity (the existing MSE loss then compares - predicted vs. ground-truth deltas, no loss-file change required); - * inference: the sampled velocity is integrated back with ``from_delta`` (cumsum) - before being returned as the emotion sequence, so the rest of the pipeline - (post-processing, metrics, rendering) sees ordinary emotion values. - -``from_delta`` is applied per generated chunk (full sequence for the offline task, -per 30-frame window for online), each anchored by its own learned first frame. -""" - -import torch -from einops import rearrange - -from framework.motion_diffusion.diffusion.matchers import LatentMatcher, DecoderLatentMatcher -from framework.motion_diffusion.diffusion.diffusion_decoder.transformer_denoiser import lengths_to_mask -from framework.motion_diffusion.diffusion.velocity_transform import to_delta, from_delta - - -class VelocityDecoderLatentMatcher(DecoderLatentMatcher): - """Decoder matcher whose diffusion data space is the frame-delta sequence. - - Identical to :class:`DecoderLatentMatcher` except for the two velocity - transforms marked below; carries no extra state, so the parent can be - promoted to this class in-place (see :class:`VelocityLatentMatcher`). - """ - - def _forward( - self, - speaker_audio_input=None, - speaker_emotion_input=None, - speaker_3dmm_input=None, - listener_emotion_input=None, - past_listener_emotion=None, - motion_length=None, - ): - with torch.no_grad(): - s_audio_encodings = self.audio_encoder._encode(speaker_audio_input) - s_audio_encodings = s_audio_encodings.repeat_interleave(self.num_preds, dim=0) - - # freeze latent RNN_VAE embedder to extract speaker latent embedding - s_latent_embed = self.latent_embedder.encode(speaker_emotion_input).unsqueeze(1) - s_latent_embed = s_latent_embed.repeat_interleave(self.num_preds, dim=0) - - s_3dmm_encodings = speaker_3dmm_input.repeat_interleave(self.num_preds, dim=0) - - s_emotion_encodings = speaker_emotion_input.repeat_interleave(self.num_preds, dim=0) - - past_listener_emotion = past_listener_emotion.repeat_interleave( - self.num_preds, dim=0) if past_listener_emotion is not None else None - - motion_length = motion_length.repeat_interleave( - self.num_preds, dim=0) if motion_length is not None else None - - model_kwargs = { - "speaker_audio_encodings": s_audio_encodings, - "speaker_latent_embed": s_latent_embed, - "speaker_3dmm_encodings": s_3dmm_encodings, - "speaker_emotion_encodings": s_emotion_encodings, - "past_listener_emotion": past_listener_emotion, - "motion_length": motion_length, - } - - if self.stage == "test": - bs, l, _ = s_audio_encodings.shape # bz * num_preds - with torch.no_grad(): - output = [output for output in self.decoder_diffusion.ddim_sample_loop_progressive( - matcher=self, - model=self.model, - model_kwargs=model_kwargs, - shape=(bs, self.window_size if self.task == "online" else l, self.emotion_dim), - )][-1] # get last output - - # --- velocity -> emotion: the diffusion sampled deltas; integrate over time. --- - output_listener_emotion = from_delta(output["sample_enc"]) # (bz * num_preds, w, d=25) - output_listener_emotion = rearrange(output_listener_emotion, - "(b n) w d -> b n w d", n=self.num_preds) - output_whole = {"prediction_emotion": output_listener_emotion} - - else: - listener_emotion_input = listener_emotion_input.repeat_interleave(self.num_preds, dim=0) - # --- emotion -> velocity: supervise the diffusion on frame deltas. --- - x_start_selected = to_delta(listener_emotion_input) # (bs * num_preds, l_w, ...) - - t, _ = self.schedule_sampler.sample(x_start_selected.shape[0], x_start_selected.device) - timesteps = t.long() - - output_whole = self.decoder_diffusion.denoise(self.model, x_start_selected, timesteps, - model_kwargs=model_kwargs) - if motion_length is not None: # offline task zero masking - device = x_start_selected.get_device() - output_mask = lengths_to_mask(motion_length, device=device, max_len=x_start_selected.shape[1]) - output_whole["prediction_emotion"] = (output_whole["prediction_emotion"] - * output_mask.float().unsqueeze(-1)) - - output_whole = {k: v.view(-1, self.num_preds, *output_whole[k].shape[1:]) for k, v in output_whole.items()} - return output_whole - - -class VelocityLatentMatcher(LatentMatcher): - """``LatentMatcher`` whose decoder diffusion runs in velocity (frame-delta) space.""" - - def __init__(self, **kwargs): - super().__init__(**kwargs) - # The velocity decoder only overrides ``_forward`` and introduces no extra - # state, so promoting the already-built (and, for test/resume, already - # checkpoint-loaded) decoder matcher in place is safe and keeps the base - # init — including weight loading — untouched. - self.diffusion_decoder.__class__ = VelocityDecoderLatentMatcher diff --git a/personalised/code/framework/motion_diffusion/diffusion/operator/__init__.py b/personalised/code/framework/motion_diffusion/diffusion/operator/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/personalised/code/framework/motion_diffusion/diffusion/operator/cross_attention.py b/personalised/code/framework/motion_diffusion/diffusion/operator/cross_attention.py deleted file mode 100644 index e7b6e777f1945451e0c65b8b949c18908a00192f..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/operator/cross_attention.py +++ /dev/null @@ -1,416 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved -""" -DETR Transformer class. -Copy-paste from torch.nn.Transformer with modifications: - * positional encodings are passed in MHattention - * extra LN at the end of encoder is removed - * decoder returns a stack of activations from all decoding layers -""" -import copy -from typing import List, Optional -from numpy import block - -import torch -import torch.nn.functional as F -from torch import Tensor, nn - - -class SkipTransformerEncoder(nn.Module): - def __init__(self, encoder_layer, num_layers, norm=None): - super().__init__() - self.d_model = encoder_layer.d_model - - self.num_layers = num_layers - self.norm = norm - - assert num_layers % 2 == 1 - - num_block = (num_layers - 1) // 2 - self.input_blocks = _get_clones(encoder_layer, num_block) - self.middle_block = _get_clone(encoder_layer) - self.output_blocks = _get_clones(encoder_layer, num_block) - self.linear_blocks = _get_clones(nn.Linear(2 * self.d_model, self.d_model), num_block) - - self._reset_parameters() - - def _reset_parameters(self): - for p in self.parameters(): - if p.dim() > 1: - nn.init.xavier_uniform_(p) - - def forward(self, src, - mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None): - x = src - - xs = [] - for module in self.input_blocks: - x = module(x, src_mask=mask, - src_key_padding_mask=src_key_padding_mask, pos=pos) - xs.append(x) - - x = self.middle_block(x, src_mask=mask, - src_key_padding_mask=src_key_padding_mask, pos=pos) - - for (module, linear) in zip(self.output_blocks, self.linear_blocks): - x = torch.cat([x, xs.pop()], dim=-1) - x = linear(x) - x = module(x, src_mask=mask, - src_key_padding_mask=src_key_padding_mask, pos=pos) - - if self.norm is not None: - x = self.norm(x) - return x - - -class SkipTransformerDecoder(nn.Module): - def __init__(self, decoder_layer, num_layers, norm=None): - super().__init__() - self.d_model = decoder_layer.d_model - - self.num_layers = num_layers - self.norm = norm - - assert num_layers % 2 == 1 - - num_block = (num_layers - 1) // 2 - self.input_blocks = _get_clones(decoder_layer, num_block) - self.middle_block = _get_clone(decoder_layer) - self.output_blocks = _get_clones(decoder_layer, num_block) - self.linear_blocks = _get_clones(nn.Linear(2 * self.d_model, self.d_model), num_block) - - self._reset_parameters() - - def _reset_parameters(self): - for p in self.parameters(): - if p.dim() > 1: - nn.init.xavier_uniform_(p) - - def forward(self, tgt, memory, - tgt_mask: Optional[Tensor] = None, - memory_mask: Optional[Tensor] = None, - tgt_key_padding_mask: Optional[Tensor] = None, - memory_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None, - query_pos: Optional[Tensor] = None): - x = tgt - - xs = [] - for module in self.input_blocks: - x = module(x, memory, tgt_mask=tgt_mask, - memory_mask=memory_mask, - tgt_key_padding_mask=tgt_key_padding_mask, - memory_key_padding_mask=memory_key_padding_mask, - pos=pos, query_pos=query_pos) - xs.append(x) - - x = self.middle_block(x, memory, tgt_mask=tgt_mask, - memory_mask=memory_mask, - tgt_key_padding_mask=tgt_key_padding_mask, - memory_key_padding_mask=memory_key_padding_mask, - pos=pos, query_pos=query_pos) - - for (module, linear) in zip(self.output_blocks, self.linear_blocks): - x = torch.cat([x, xs.pop()], dim=-1) - x = linear(x) - x = module(x, memory, tgt_mask=tgt_mask, - memory_mask=memory_mask, - tgt_key_padding_mask=tgt_key_padding_mask, - memory_key_padding_mask=memory_key_padding_mask, - pos=pos, query_pos=query_pos) - - if self.norm is not None: - x = self.norm(x) - - return x - - -class Transformer(nn.Module): - - def __init__(self, d_model=512, nhead=8, num_encoder_layers=6, - num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, - activation="relu", normalize_before=False, - return_intermediate_dec=False): - super().__init__() - - encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, - dropout, activation, normalize_before) - encoder_norm = nn.LayerNorm(d_model) if normalize_before else None - self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) - - decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, - dropout, activation, normalize_before) - decoder_norm = nn.LayerNorm(d_model) - self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm, - return_intermediate=return_intermediate_dec) - - self._reset_parameters() - - self.d_model = d_model - self.nhead = nhead - - def _reset_parameters(self): - for p in self.parameters(): - if p.dim() > 1: - nn.init.xavier_uniform_(p) - - def forward(self, src, mask, query_embed, pos_embed): - # flatten NxCxHxW to HWxNxC - bs, c, h, w = src.shape - src = src.flatten(2).permute(2, 0, 1) - pos_embed = pos_embed.flatten(2).permute(2, 0, 1) - query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1) - mask = mask.flatten(1) - - tgt = torch.zeros_like(query_embed) - memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed) - hs = self.decoder(tgt, memory, memory_key_padding_mask=mask, - pos=pos_embed, query_pos=query_embed) - return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w) - - -class TransformerEncoder(nn.Module): - - def __init__(self, encoder_layer, num_layers, norm=None): - super().__init__() - self.layers = _get_clones(encoder_layer, num_layers) - self.num_layers = num_layers - self.norm = norm - - def forward(self, src, - mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None): - output = src - - for layer in self.layers: - output = layer(output, src_mask=mask, - src_key_padding_mask=src_key_padding_mask, pos=pos) - - if self.norm is not None: - output = self.norm(output) - - return output - - -class TransformerDecoder(nn.Module): - - def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False): - super().__init__() - self.layers = _get_clones(decoder_layer, num_layers) - self.num_layers = num_layers - self.norm = norm - self.return_intermediate = return_intermediate - - def forward(self, tgt, memory, - tgt_mask: Optional[Tensor] = None, - memory_mask: Optional[Tensor] = None, - tgt_key_padding_mask: Optional[Tensor] = None, - memory_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None, - query_pos: Optional[Tensor] = None): - output = tgt - - intermediate = [] - - for layer in self.layers: - output = layer(output, memory, tgt_mask=tgt_mask, - memory_mask=memory_mask, - tgt_key_padding_mask=tgt_key_padding_mask, - memory_key_padding_mask=memory_key_padding_mask, - pos=pos, query_pos=query_pos) - if self.return_intermediate: - intermediate.append(self.norm(output)) - - if self.norm is not None: - output = self.norm(output) - if self.return_intermediate: - intermediate.pop() - intermediate.append(output) - - if self.return_intermediate: - return torch.stack(intermediate) - - return output.unsqueeze(0) - - -class TransformerEncoderLayer(nn.Module): - - def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, - activation="relu", normalize_before=False): - super().__init__() - self.d_model = d_model - self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) - # Implementation of Feedforward model - self.linear1 = nn.Linear(d_model, dim_feedforward) - self.dropout = nn.Dropout(dropout) - self.linear2 = nn.Linear(dim_feedforward, d_model) - - self.norm1 = nn.LayerNorm(d_model) - self.norm2 = nn.LayerNorm(d_model) - self.dropout1 = nn.Dropout(dropout) - self.dropout2 = nn.Dropout(dropout) - - self.activation = _get_activation_fn(activation) - self.normalize_before = normalize_before - - def with_pos_embed(self, tensor, pos: Optional[Tensor]): - return tensor if pos is None else tensor + pos - - def forward_post(self, - src, - src_mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None): - q = k = self.with_pos_embed(src, pos) - src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, - key_padding_mask=src_key_padding_mask)[0] - src = src + self.dropout1(src2) - src = self.norm1(src) - src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) - src = src + self.dropout2(src2) - src = self.norm2(src) - return src - - def forward_pre(self, src, - src_mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None): - src2 = self.norm1(src) - q = k = self.with_pos_embed(src2, pos) - src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask, - key_padding_mask=src_key_padding_mask)[0] - src = src + self.dropout1(src2) - src2 = self.norm2(src) - src2 = self.linear2(self.dropout(self.activation(self.linear1(src2)))) - src = src + self.dropout2(src2) - return src - - def forward(self, src, - src_mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None): - if self.normalize_before: - return self.forward_pre(src, src_mask, src_key_padding_mask, pos) - return self.forward_post(src, src_mask, src_key_padding_mask, pos) - - -class TransformerDecoderLayer(nn.Module): - - def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, - activation="relu", normalize_before=False): - super().__init__() - self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) - self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) - # Implementation of Feedforward model - self.d_model = d_model - self.linear1 = nn.Linear(d_model, dim_feedforward) - self.dropout = nn.Dropout(dropout) - self.linear2 = nn.Linear(dim_feedforward, d_model) - - self.norm1 = nn.LayerNorm(d_model) - self.norm2 = nn.LayerNorm(d_model) - self.norm3 = nn.LayerNorm(d_model) - self.dropout1 = nn.Dropout(dropout) - self.dropout2 = nn.Dropout(dropout) - self.dropout3 = nn.Dropout(dropout) - - self.activation = _get_activation_fn(activation) - self.normalize_before = normalize_before - - def with_pos_embed(self, tensor, pos: Optional[Tensor]): - return tensor if pos is None else tensor + pos - - def forward_post(self, tgt, memory, - tgt_mask: Optional[Tensor] = None, - memory_mask: Optional[Tensor] = None, - tgt_key_padding_mask: Optional[Tensor] = None, - memory_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None, - query_pos: Optional[Tensor] = None): - q = k = self.with_pos_embed(tgt, query_pos) - tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, - key_padding_mask=tgt_key_padding_mask)[0] - - tgt = tgt + self.dropout1(tgt2) - tgt = self.norm1(tgt) - tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos), - key=self.with_pos_embed(memory, pos), - value=memory, attn_mask=memory_mask, - key_padding_mask=memory_key_padding_mask)[0] - - tgt = tgt + self.dropout2(tgt2) - tgt = self.norm2(tgt) - tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) - tgt = tgt + self.dropout3(tgt2) - tgt = self.norm3(tgt) - return tgt - - def forward_pre(self, tgt, memory, - tgt_mask: Optional[Tensor] = None, - memory_mask: Optional[Tensor] = None, - tgt_key_padding_mask: Optional[Tensor] = None, - memory_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None, - query_pos: Optional[Tensor] = None): - tgt2 = self.norm1(tgt) - q = k = self.with_pos_embed(tgt2, query_pos) - tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, - key_padding_mask=tgt_key_padding_mask)[0] - tgt = tgt + self.dropout1(tgt2) - tgt2 = self.norm2(tgt) - tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos), - key=self.with_pos_embed(memory, pos), - value=memory, attn_mask=memory_mask, - key_padding_mask=memory_key_padding_mask)[0] - tgt = tgt + self.dropout2(tgt2) - tgt2 = self.norm3(tgt) - tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) - tgt = tgt + self.dropout3(tgt2) - return tgt - - def forward(self, tgt, memory, - tgt_mask: Optional[Tensor] = None, - memory_mask: Optional[Tensor] = None, - tgt_key_padding_mask: Optional[Tensor] = None, - memory_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None, - query_pos: Optional[Tensor] = None): - if self.normalize_before: - return self.forward_pre(tgt, memory, tgt_mask, memory_mask, - tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) - return self.forward_post(tgt, memory, tgt_mask, memory_mask, - tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) - - -def _get_clone(module): - return copy.deepcopy(module) - - -def _get_clones(module, N): - return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) - - -def build_transformer(args): - return Transformer( - d_model=args.hidden_dim, - dropout=args.dropout, - nhead=args.nheads, - dim_feedforward=args.dim_feedforward, - num_encoder_layers=args.enc_layers, - num_decoder_layers=args.dec_layers, - normalize_before=args.pre_norm, - return_intermediate_dec=True, - ) - - -def _get_activation_fn(activation): - """Return an activation function given a string""" - if activation == "relu": - return F.relu - if activation == "gelu": - return F.gelu - if activation == "glu": - return F.glu - raise RuntimeError(F"activation should be relu/gelu, not {activation}.") diff --git a/personalised/code/framework/motion_diffusion/diffusion/operator/embeddings.py b/personalised/code/framework/motion_diffusion/diffusion/operator/embeddings.py deleted file mode 100644 index 179582bd988bc996c2d056f739af298c1da9fd06..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/operator/embeddings.py +++ /dev/null @@ -1,320 +0,0 @@ -# This file is taken from signjoey repository -import math -import torch -from torch import Tensor, nn - - -def get_activation(activation_type): - if activation_type == "relu": - return nn.ReLU() - elif activation_type == "relu6": - return nn.ReLU6() - elif activation_type == "prelu": - return nn.PReLU() - elif activation_type == "selu": - return nn.SELU() - elif activation_type == "celu": - return nn.CELU() - elif activation_type == "gelu": - return nn.GELU() - elif activation_type == "sigmoid": - return nn.Sigmoid() - elif activation_type == "softplus": - return nn.Softplus() - elif activation_type == "softshrink": - return nn.Softshrink() - elif activation_type == "softsign": - return nn.Softsign() - elif activation_type == "tanh": - return nn.Tanh() - elif activation_type == "tanhshrink": - return nn.Tanhshrink() - else: - raise ValueError("Unknown activation type {}".format(activation_type)) - - -class MaskedNorm(nn.Module): - """ - Original Code from: - https://discuss.pytorch.org/t/batchnorm-for-different-sized-samples-in-batch/44251/8 - """ - - def __init__(self, norm_type, num_groups, num_features): - super().__init__() - self.norm_type = norm_type - if self.norm_type == "batch": - self.norm = nn.BatchNorm1d(num_features=num_features) - elif self.norm_type == "group": - self.norm = nn.GroupNorm(num_groups=num_groups, num_channels=num_features) - elif self.norm_type == "layer": - self.norm = nn.LayerNorm(normalized_shape=num_features) - else: - raise ValueError("Unsupported Normalization Layer") - - self.num_features = num_features - - def forward(self, x: Tensor, mask: Tensor): - if self.training: - reshaped = x.reshape([-1, self.num_features]) - reshaped_mask = mask.reshape([-1, 1]) > 0 - selected = torch.masked_select(reshaped, reshaped_mask).reshape( - [-1, self.num_features] - ) - batch_normed = self.norm(selected) - scattered = reshaped.masked_scatter(reshaped_mask, batch_normed) - return scattered.reshape([x.shape[0], -1, self.num_features]) - else: - reshaped = x.reshape([-1, self.num_features]) - batched_normed = self.norm(reshaped) - return batched_normed.reshape([x.shape[0], -1, self.num_features]) - - -# TODO (Cihan): Spatial and Word Embeddings are pretty much the same -# We might as well convert them into a single module class. -# Only difference is the lut vs linear layers. -class Embeddings(nn.Module): - """ - Simple embeddings class - """ - - # pylint: disable=unused-argument - def __init__( - self, - embedding_dim: int = 64, - num_heads: int = 8, - scale: bool = False, - scale_factor: float = None, - norm_type: str = None, - activation_type: str = None, - vocab_size: int = 0, - padding_idx: int = 1, - freeze: bool = False, - **kwargs - ): - """ - Create new embeddings for the vocabulary. - Use scaling for the Transformer. - - :param embedding_dim: - :param scale: - :param vocab_size: - :param padding_idx: - :param freeze: freeze the embeddings during training - """ - super().__init__() - - self.embedding_dim = embedding_dim - self.vocab_size = vocab_size - self.lut = nn.Embedding(vocab_size, self.embedding_dim, padding_idx=padding_idx) - - self.norm_type = norm_type - if self.norm_type: - self.norm = MaskedNorm( - norm_type=norm_type, num_groups=num_heads, num_features=embedding_dim - ) - - self.activation_type = activation_type - if self.activation_type: - self.activation = get_activation(activation_type) - - self.scale = scale - if self.scale: - if scale_factor: - self.scale_factor = scale_factor - else: - self.scale_factor = math.sqrt(self.embedding_dim) - - # if freeze: - # freeze_params(self) - - # pylint: disable=arguments-differ - def forward(self, x: Tensor, mask: Tensor = None) -> Tensor: - """ - Perform lookup for input `x` in the embedding table. - - :param mask: token masks - :param x: index in the vocabulary - :return: embedded representation for `x` - """ - - x = self.lut(x) - - if self.norm_type: - x = self.norm(x, mask) - - if self.activation_type: - x = self.activation(x) - - if self.scale: - return x * self.scale_factor - else: - return x - - def __repr__(self): - return "%s(embedding_dim=%d, vocab_size=%d)" % ( - self.__class__.__name__, - self.embedding_dim, - self.vocab_size, - ) - - -class SpatialEmbeddings(nn.Module): - """ - Simple Linear Projection Layer - (For encoder outputs to predict glosses) - """ - - # pylint: disable=unused-argument - def __init__( - self, - embedding_dim: int, - input_size: int, - num_heads: int, - freeze: bool = False, - norm_type: str = "batch", - activation_type: str = "softsign", - scale: bool = False, - scale_factor: float = None, - **kwargs - ): - """ - Create new embeddings for the vocabulary. - Use scaling for the Transformer. - - :param embedding_dim: - :param input_size: - :param freeze: freeze the embeddings during training - """ - super().__init__() - - self.embedding_dim = embedding_dim - self.input_size = input_size - self.ln = nn.Linear(self.input_size, self.embedding_dim) - - self.norm_type = norm_type - if self.norm_type: - self.norm = MaskedNorm( - norm_type=norm_type, num_groups=num_heads, num_features=embedding_dim - ) - - self.activation_type = activation_type - if self.activation_type: - self.activation = get_activation(activation_type) - - self.scale = scale - if self.scale: - if scale_factor: - self.scale_factor = scale_factor - else: - self.scale_factor = math.sqrt(self.embedding_dim) - - # if freeze: - # freeze_params(self) - - # pylint: disable=arguments-differ - def forward(self, x: Tensor, mask: Tensor) -> Tensor: - """ - :param mask: frame masks - :param x: input frame features - :return: embedded representation for `x` - """ - - x = self.ln(x) - - if self.norm_type: - x = self.norm(x, mask) - - if self.activation_type: - x = self.activation(x) - - if self.scale: - return x * self.scale_factor - else: - return x - - def __repr__(self): - return "%s(embedding_dim=%d, input_size=%d)" % ( - self.__class__.__name__, - self.embedding_dim, - self.input_size, - ) - - -def get_timestep_embedding( - timesteps: torch.Tensor, - embedding_dim: int, - flip_sin_to_cos: bool = False, - downscale_freq_shift: float = 1, - scale: float = 1, - max_period: int = 10000, -): - """ - This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. - - :param timesteps: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the - embeddings. :return: an [N x dim] Tensor of positional embeddings. - """ - assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" - - half_dim = embedding_dim // 2 - exponent = -math.log(max_period) * torch.arange( - start=0, end=half_dim, dtype=torch.float32, device=timesteps.device - ) - exponent = exponent / (half_dim - downscale_freq_shift) - - emb = torch.exp(exponent) - emb = timesteps[:, None].float() * emb[None, :] - - # scale embeddings - emb = scale * emb - - # concat sine and cosine embeddings - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) - - # flip sine and cosine embeddings - if flip_sin_to_cos: - emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) - - # zero pad - if embedding_dim % 2 == 1: - emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) - return emb - - -class TimestepEmbedding(nn.Module): - def __init__(self, channel: int, time_embed_dim: int, act_fn: str = "silu"): - super().__init__() - - self.linear_1 = nn.Linear(channel, time_embed_dim) - self.act = None - if act_fn == "silu": - self.act = nn.SiLU() - self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim) - - def forward(self, sample): - sample = self.linear_1(sample) - - if self.act is not None: - sample = self.act(sample) - - sample = self.linear_2(sample) - return sample - - -class Timesteps(nn.Module): - def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float): - super().__init__() - self.num_channels = num_channels - self.flip_sin_to_cos = flip_sin_to_cos - self.downscale_freq_shift = downscale_freq_shift - - def forward(self, timesteps): - t_emb = get_timestep_embedding( - timesteps, - self.num_channels, - flip_sin_to_cos=self.flip_sin_to_cos, - downscale_freq_shift=self.downscale_freq_shift, - ) - return t_emb diff --git a/personalised/code/framework/motion_diffusion/diffusion/operator/position_encoding.py b/personalised/code/framework/motion_diffusion/diffusion/operator/position_encoding.py deleted file mode 100644 index a749559173c88ba4401ba736890393fb42fa538a..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/operator/position_encoding.py +++ /dev/null @@ -1,179 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved -""" -Various positional encodings for the transformer. -""" -import math -from typing import List, Optional -import numpy as np -import torch -from torch import Tensor, nn - - -class NestedTensor(object): - - def __init__(self, tensors, mask: Optional[Tensor]): - self.tensors = tensors - self.mask = mask - - def to(self, device): - # type: (Device) -> NestedTensor # noqa - cast_tensor = self.tensors.to(device) - mask = self.mask - if mask is not None: - assert mask is not None - cast_mask = mask.to(device) - else: - cast_mask = None - return NestedTensor(cast_tensor, cast_mask) - - def decompose(self): - return self.tensors, self.mask - - def __repr__(self): - return str(self.tensors) - - -class PositionEmbeddingSine(nn.Module): - """ - This is a more standard version of the position embedding, very similar to the one - used by the Attention is all you need paper, generalized to work on images. - """ - - def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): - super().__init__() - self.num_pos_feats = num_pos_feats - self.temperature = temperature - self.normalize = normalize - if scale is not None and normalize is False: - raise ValueError("normalize should be True if scale is passed") - if scale is None: - scale = 2 * math.pi - self.scale = scale - - def forward(self, tensor_list: NestedTensor): - x = tensor_list.tensors - mask = tensor_list.mask - assert mask is not None - not_mask = ~mask - y_embed = not_mask.cumsum(1, dtype=torch.float32) - x_embed = not_mask.cumsum(2, dtype=torch.float32) - if self.normalize: - eps = 1e-6 - y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale - x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale - - dim_t = torch.arange(self.num_pos_feats, - dtype=torch.float32, device=x.device) - dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) - - pos_x = x_embed[:, :, :, None] / dim_t - pos_y = y_embed[:, :, :, None] / dim_t - pos_x = torch.stack( - (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) - pos_y = torch.stack( - (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) - pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) - return pos - - -class PositionEmbeddingLearned(nn.Module): - """ - Absolute pos embedding, learned. - """ - - def __init__(self, num_pos_feats=256): - super().__init__() - self.row_embed = nn.Embedding(50, num_pos_feats) - self.col_embed = nn.Embedding(50, num_pos_feats) - self.reset_parameters() - - def reset_parameters(self): - nn.init.uniform_(self.row_embed.weight) - nn.init.uniform_(self.col_embed.weight) - - def forward(self, tensor_list: NestedTensor): - x = tensor_list.tensors - h, w = x.shape[-2:] - i = torch.arange(w, device=x.device) - j = torch.arange(h, device=x.device) - x_emb = self.col_embed(i) - y_emb = self.row_embed(j) - pos = torch.cat([ - x_emb.unsqueeze(0).repeat(h, 1, 1), - y_emb.unsqueeze(1).repeat(1, w, 1), - ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1) - return pos - - -class PositionEmbeddingSine1D(nn.Module): - - def __init__(self, d_model, max_len=5000, batch_first=False): - super().__init__() - self.batch_first = batch_first - - pe = torch.zeros(max_len, d_model) - position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) - div_term = torch.exp(torch.arange( - 0, d_model, 2).float() * (-np.log(10000.0) / d_model)) - pe[:, 0::2] = torch.sin(position * div_term) - pe[:, 1::2] = torch.cos(position * div_term) - pe = pe.unsqueeze(0).transpose(0, 1) - - self.register_buffer('pe', pe) - - def forward(self, x): - # not used in the final model - if self.batch_first: - pos = self.pe.permute(1, 0, 2)[:, :x.shape[1], :] - else: - pos = self.pe[:x.shape[0], :] - return pos - - -class PositionEmbeddingLearned1D(nn.Module): - - def __init__(self, d_model, max_len=5000, batch_first=False): - super().__init__() - self.batch_first = batch_first - # self.dropout = nn.Dropout(p=dropout) - - self.pe = nn.Parameter(torch.zeros(max_len, 1, d_model)) - - self.reset_parameters() - - def reset_parameters(self): - nn.init.uniform_(self.pe) - - def forward(self, x): - # not used in the final model - if self.batch_first: - pos = self.pe.permute(1, 0, 2)[:, :x.shape[1], :] - else: - x = x + self.pe[:x.shape[0], :] - return x - # return self.dropout(x) - - -def build_position_encoding(N_steps, - position_embedding="sine", - embedding_dim="1D"): - # N_steps = hidden_dim // 2 - if embedding_dim == "1D": - if position_embedding in ('v2', 'sine'): - position_embedding = PositionEmbeddingSine1D(N_steps) - elif position_embedding in ('v3', 'learned'): - position_embedding = PositionEmbeddingLearned1D(N_steps) - else: - raise ValueError(f"not supported {position_embedding}") - elif embedding_dim == "2D": - if position_embedding in ('v2', 'sine'): - # TODO find a better way of exposing other arguments - position_embedding = PositionEmbeddingSine(N_steps, normalize=True) - elif position_embedding in ('v3', 'learned'): - position_embedding = PositionEmbeddingLearned(N_steps) - else: - raise ValueError(f"not supported {position_embedding}") - else: - raise ValueError(f"not supported {embedding_dim}") - - return position_embedding diff --git a/personalised/code/framework/motion_diffusion/diffusion/resample.py b/personalised/code/framework/motion_diffusion/diffusion/resample.py deleted file mode 100644 index c067cfdb09326c7fc2b399da707695c556af3334..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/resample.py +++ /dev/null @@ -1,153 +0,0 @@ -from abc import ABC, abstractmethod -import numpy as np -import torch as th -import torch.distributed as dist - - -def create_named_schedule_sampler(name, diffusion): - """ - Create a ScheduleSampler from a library of pre-defined samplers. - - :param name: the name of the sampler. - :param diffusion: the diffusion object to sample for. - """ - if name == "uniform": - return UniformSampler(diffusion) - elif name == "loss-second-moment": - return LossSecondMomentResampler(diffusion) - else: - raise NotImplementedError(f"unknown schedule sampler: {name}") - - -class ScheduleSampler(ABC): - """ - A distribution over timesteps in the diffusion process, intended to reduce - variance of the objective. - - By default, samplers perform unbiased importance sampling, in which the - objective's mean is unchanged. - However, subclasses may override sample() to change how the resampled - terms are reweighted, allowing for actual changes in the objective. - """ - - @abstractmethod - def weights(self): - """ - Get a numpy array of weights, one per diffusion step. - - The weights needn't be normalized, but must be positive. - """ - - def sample(self, batch_size, device): - """ - Importance-sample timesteps for a batch. - - :param batch_size: the number of timesteps. - :param device: the torch device to save to. - :return: a tuple (timesteps, weights): - - timesteps: a tensor of timestep indices. - - weights: a tensor of weights to scale the resulting losses. - """ - w = self.weights() - p = w / np.sum(w) - indices_np = np.random.choice(len(p), size=(batch_size,), p=p) - indices = th.from_numpy(indices_np).long().to(device) - weights_np = 1 / (len(p) * p[indices_np]) - weights = th.from_numpy(weights_np).float().to(device) - return indices, weights - - -class UniformSampler(ScheduleSampler): - def __init__(self, diffusion): - self.diffusion = diffusion - self._weights = np.ones([self.diffusion.num_timesteps]) # shape: (num_timesteps, ) - - def weights(self): - return self._weights - - -class LossAwareSampler(ScheduleSampler): - def update_with_local_losses(self, local_ts, local_losses): - """ - Update the reweighting using losses from a model. - - Call this method from each rank with a batch of timesteps and the - corresponding losses for each of those timesteps. - This method will perform synchronization to make sure all of the ranks - maintain the exact same reweighting. - - :param local_ts: an integer Tensor of timesteps. - :param local_losses: a 1D Tensor of losses. - """ - batch_sizes = [ - th.tensor([0], dtype=th.int32, device=local_ts.device) - for _ in range(dist.get_world_size()) - ] - dist.all_gather( - batch_sizes, - th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), - ) - - # Pad all_gather batches to be the maximum batch size. - batch_sizes = [x.item() for x in batch_sizes] - max_bs = max(batch_sizes) - - timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] - loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] - dist.all_gather(timestep_batches, local_ts) - dist.all_gather(loss_batches, local_losses) - timesteps = [ - x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] - ] - losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] - self.update_with_all_losses(timesteps, losses) - - @abstractmethod - def update_with_all_losses(self, ts, losses): - """ - Update the reweighting using losses from a model. - - Sub-classes should override this method to update the reweighting - using losses from the model. - - This method directly updates the reweighting without synchronizing - between workers. It is called by update_with_local_losses from all - ranks with identical arguments. Thus, it should have deterministic - behavior to maintain state across workers. - - :param ts: a list of int timesteps. - :param losses: a list of float losses, one per timestep. - """ - - -class LossSecondMomentResampler(LossAwareSampler): - def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): - self.diffusion = diffusion - self.history_per_term = history_per_term - self.uniform_prob = uniform_prob - self._loss_history = np.zeros( - [diffusion.num_timesteps, history_per_term], dtype=np.float64 - ) - self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) - - def weights(self): - if not self._warmed_up(): - return np.ones([self.diffusion.num_timesteps], dtype=np.float64) - weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1)) - weights /= np.sum(weights) - weights *= 1 - self.uniform_prob - weights += self.uniform_prob / len(weights) - return weights - - def update_with_all_losses(self, ts, losses): - for t, loss in zip(ts, losses): - if self._loss_counts[t] == self.history_per_term: - # Shift out the oldest loss term. - self._loss_history[t, :-1] = self._loss_history[t, 1:] - self._loss_history[t, -1] = loss - else: - self._loss_history[t, self._loss_counts[t]] = loss - self._loss_counts[t] += 1 - - def _warmed_up(self): - return (self._loss_counts == self.history_per_term).all() diff --git a/personalised/code/framework/motion_diffusion/diffusion/rnn.py b/personalised/code/framework/motion_diffusion/diffusion/rnn.py deleted file mode 100644 index 53532c2141f9d3288b9077acc2b8a91c890954b9..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/rnn.py +++ /dev/null @@ -1,282 +0,0 @@ -import os -from pathlib import Path -from typing import List -import hydra -import numpy as np -import torch -import torch.nn as nn -from einops import rearrange -from framework.motion_diffusion.diffusion.torch import init_weights, zeros, batch_to -from framework.utils.util import save_checkpoint - -nl = { - "tanh": nn.Tanh, - "relu": nn.ReLU, - "sigmoid": nn.Sigmoid, - "elu": nn.ELU, - "selu": nn.SELU, - "softplus": nn.Softplus, - "softsign": nn.Softsign, - "leaky_relu": nn.LeakyReLU, - "none": lambda x: x, -} - - -def sample(mu): - return torch.randn_like(mu) - - -def rc(x_start, pred, batch_first=True): - # x_start -> [batch_size, ...] - # pred -> [seq_length, batch_size, ...] | [batch_size, seq_length, ...] - if batch_first: - x_start = x_start.unsqueeze(1) - shapes = [1 for s in x_start.shape] - shapes[1] = pred.shape[1] - x_start = x_start.repeat(shapes) - else: - x_start = x_start.unsqueeze(0) - shapes = [1 for s in x_start.shape] - shapes[0] = pred.shape[0] - x_start = x_start.repeat(shapes) - return x_start + pred - - -def rc_recurrent(x_start, pred, batch_first=True): # residual connection => offsets modeling - # x_start -> [batch_size, ...] - # pred -> [seq_length, batch_size, ...] | [batch_size, seq_length, ...] - if batch_first: - pred[:, 0] = x_start + pred[:, 0] - for i in range(1, pred.shape[1]): - pred[:, i] = pred[:, i - 1] + pred[:, i] - else: # seq length first - pred[0] = x_start + pred[0] - for i in range(1, pred.shape[0]): - pred[i] = pred[i - 1] + pred[i] - return pred - - -class BasicMLP(nn.Module): - def __init__(self, input_dim, output_dim, hidden_dims=[], dropout=0.5, non_linearities='relu'): - super(BasicMLP, self).__init__() - self.non_linearities = non_linearities - - self.dropout = nn.Dropout(dropout) - self.nl = nl[non_linearities]() - - self.denses = None - - # hidden dims - hidden_dims = hidden_dims + [output_dim, ] # output dim is treated as the last hidden dim - - seqs = [] - for i in range(len(hidden_dims)): - linear = nn.Linear(input_dim if i == 0 else hidden_dims[i - 1], hidden_dims[i]) - init_weights(linear) - seqs.append(nn.Sequential(self.dropout, linear, self.nl)) - - self.denses = nn.Sequential(*seqs) - - def forward(self, x): - return self.denses(x) if self.denses is not None else x - - -class MLP(nn.Module): - # https://github.com/Khrylx/DLow - def __init__(self, input_dim, hidden_dims=(128, 128), activation='tanh'): - super(MLP, self).__init__() - if activation == 'tanh': - self.activation = torch.tanh - elif activation == 'relu': - self.activation = torch.relu - elif activation == 'sigmoid': - self.activation = torch.sigmoid - - self.out_dim = hidden_dims[-1] - self.affine_layers = nn.ModuleList() - last_dim = input_dim - for nh in hidden_dims: - self.affine_layers.append(nn.Linear(last_dim, nh)) - last_dim = nh - - def forward(self, x): - for affine in self.affine_layers: - x = self.activation(affine(x)) - return x - - -class RNN(nn.Module): - # https://github.com/Khrylx/DLow - def __init__(self, input_dim, out_dim, cell_type='lstm', bi_dir=False): - super(RNN, self).__init__() - self.input_dim = input_dim - self.out_dim = out_dim - self.cell_type = cell_type - self.bi_dir = bi_dir - self.mode = 'batch' - rnn_cls = nn.LSTMCell if cell_type == 'lstm' else nn.GRUCell - hidden_dim = out_dim // 2 if bi_dir else out_dim - self.rnn_f = rnn_cls(self.input_dim, hidden_dim) - if bi_dir: - self.rnn_b = rnn_cls(self.input_dim, hidden_dim) - self.hx, self.cx = None, None - - def set_mode(self, mode): - self.mode = mode - - def initialize(self, batch_size=1, hx=None, cx=None): - if self.mode == 'step': - self.hx = zeros((batch_size, self.rnn_f.hidden_size)) if hx is None else hx - if self.cell_type == 'lstm': - self.cx = zeros((batch_size, self.rnn_f.hidden_size)) if cx is None else cx - - def forward(self, x): - if self.mode == 'step': - self.hx, self.cx = batch_to(x.device, self.hx, self.cx) - if self.cell_type == 'lstm': - self.hx, self.cx = self.rnn_f(x, (self.hx, self.cx)) - else: - self.hx = self.rnn_f(x, self.hx) - rnn_out = self.hx - else: - rnn_out_f = self.batch_forward(x) - if not self.bi_dir: - return rnn_out_f - rnn_out_b = self.batch_forward(x, reverse=True) - rnn_out = torch.cat((rnn_out_f, rnn_out_b), 2) - return rnn_out - - def batch_forward(self, x, reverse=False): - rnn = self.rnn_b if reverse else self.rnn_f - rnn_out = [] - hx = zeros((x.size(1), rnn.hidden_size), device=x.device) - if self.cell_type == 'lstm': - cx = zeros((x.size(1), rnn.hidden_size), device=x.device) - ind = reversed(range(x.size(0))) if reverse else range(x.size(0)) - for t in ind: - if self.cell_type == 'lstm': - hx, cx = rnn(x[t, ...], (hx, cx)) - else: - hx = rnn(x[t, ...], hx) - rnn_out.append(hx.unsqueeze(0)) - if reverse: - rnn_out.reverse() - rnn_out = torch.cat(rnn_out, 0) - return rnn_out - - -class LatentEmbedder(nn.Module): - def __init__(self, - hidden_dim: int = 512, - z_dim: int = 512, - emb_dims: List[int] = [128, 128], - num_layers: int = 2, - rnn_type: str = 'gru', - dropout: float = 0.0, - emotion_dim: int = 25, - coeff_3dmm_dim: int = 58, - **kwargs, - ): - super(LatentEmbedder, self).__init__() - self.hidden_dim = hidden_dim - self.z_dim = z_dim - self.emb_dims = emb_dims - self.num_layers = num_layers - self.rnn_type = rnn_type - self.dropout = dropout - self.emotion_dim = emotion_dim - self.coeff_3dmm_dim = coeff_3dmm_dim - - # encode - self.x_rnn = RNN(self.emotion_dim, self.hidden_dim, cell_type=self.rnn_type) - # z - self.fc_mu_enc = nn.Linear(self.hidden_dim, self.z_dim) - self.fc_logvar_enc = nn.Linear(self.hidden_dim, self.z_dim) - - # decode - self.fc_z_dec = nn.Linear(self.z_dim, self.hidden_dim) - self.d_rnn = RNN(self.emotion_dim + self.hidden_dim, self.hidden_dim, cell_type=self.rnn_type) - self.d_mlp = MLP(self.hidden_dim, self.emb_dims) - self.d_out = nn.Linear(self.d_mlp.out_dim, self.emotion_dim) - self.d_rnn.set_mode('step') - # decode_3dmm - # linear layers from emotion_dim to coeff_3dmm_dim - self.coeff_reg = nn.Sequential( - nn.Linear(self.emotion_dim, self.hidden_dim), - nn.Tanh(), - nn.Linear(self.hidden_dim, self.coeff_3dmm_dim), - ) - - self.dropout = nn.Dropout(self.dropout) - - def _encode(self, x): - return self.x_rnn(x)[-1] - - def get_encodings(self, emotion_seq): - emotion_seq = rearrange(emotion_seq, 'b s f -> s b f') - emotion_seq = self.x_rnn(emotion_seq) - return rearrange(emotion_seq, 's b f -> b s f') - - def _decode(self, h_y, seq_len=None): - h_y = self.fc_z_dec(h_y) - h_y = self.dropout(h_y) - self.d_rnn.initialize(batch_size=h_y.shape[0]) - y = [] - for i in range(seq_len): - y_p = torch.zeros((h_y.shape[0], self.emotion_dim), device=h_y.device) if i == 0 else y_i - rnn_in = torch.cat([h_y, y_p], dim=1) - h = self.d_rnn(rnn_in) - h = self.d_mlp(h) - y_i = self.d_out(h) - y.append(y_i) - - return torch.stack(y) - - def encode(self, emotion_seq): - return self._encode(rearrange(emotion_seq, 'b s f -> s b f')) - - def decode(self, emb): - Y_r = self._decode(emb) - - # BACK TO ORIGINAL SHAPE - Y_r = rearrange(Y_r, "s b f -> b s f") - return Y_r - - def decode_coeff(self, emotion): - return self.coeff_reg(emotion) - - def reparameterize(self, mu, logvar): - std = torch.exp(0.5 * logvar) - eps = torch.randn_like(std) - return mu + eps * std # 128 + 128 * 128 - - def forward(self, emotion=None, _3dmm=None, **kwargs): - target_emotion = emotion - - emotion = rearrange(emotion, 'b s f -> s b f') - h_x = self._encode(emotion) # h_x.shape: (bz, dim) - mu = self.fc_mu_enc(h_x) # mu.shape: (bz, dim) - logvar = self.fc_logvar_enc(h_x) # logvar.bz: (bz, dim) - - if self.training: - z = self.reparameterize(mu, logvar) - else: - z = mu - - Y_r = self._decode(z, seq_len=emotion.shape[0]) # (seq_len, bz, 25) - Y_r = rearrange(Y_r, "s b f -> b s f") - - return { - "prediction": Y_r, - "target": target_emotion, # we are autoencoding - "coefficients_3dmm": self.coeff_reg(Y_r), - "target_coefficients": _3dmm, - "mu": mu, - "logvar": logvar, - } - - def save_ckpt(self, ckpt_path, optimizer): - save_checkpoint(ckpt_path, self, optimizer) - - def get_model_name(self): - return self.__class__.__name__ \ No newline at end of file diff --git a/personalised/code/framework/motion_diffusion/diffusion/torch.py b/personalised/code/framework/motion_diffusion/diffusion/torch.py deleted file mode 100644 index d1bcd5c5fc0570d5eadfbeb10473b262f33b6f03..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/torch.py +++ /dev/null @@ -1,242 +0,0 @@ -import torch -import numpy as np -from torch.optim import lr_scheduler -import torch.nn as nn - -# https://raw.githubusercontent.com/Khrylx/DLow/master/utils/torch.py - -tensor = torch.tensor -DoubleTensor = torch.DoubleTensor -FloatTensor = torch.FloatTensor -LongTensor = torch.LongTensor -ByteTensor = torch.ByteTensor -ones = torch.ones -zeros = torch.zeros - - -class to_cpu: - - def __init__(self, *models): - self.models = list(filter(lambda x: x is not None, models)) - self.prev_devices = [x.device if hasattr(x, 'device') else next(x.parameters()).device for x in self.models] - for x in self.models: - x.to(torch.device('cpu')) - - def __enter__(self): - pass - - def __exit__(self, *args): - for x, device in zip(self.models, self.prev_devices): - x.to(device) - return False - - -class to_device: - - def __init__(self, device, *models): - self.models = list(filter(lambda x: x is not None, models)) - self.prev_devices = [x.device if hasattr(x, 'device') else next(x.parameters()).device for x in self.models] - for x in self.models: - x.to(device) - - def __enter__(self): - pass - - def __exit__(self, *args): - for x, device in zip(self.models, self.prev_devices): - x.to(device) - return False - - -class to_test: - - def __init__(self, *models): - self.models = list(filter(lambda x: x is not None, models)) - self.prev_modes = [x.training for x in self.models] - for x in self.models: - x.train(False) - - def __enter__(self): - pass - - def __exit__(self, *args): - for x, mode in zip(self.models, self.prev_modes): - x.train(mode) - return False - - -class to_train: - - def __init__(self, *models): - self.models = list(filter(lambda x: x is not None, models)) - self.prev_modes = [x.training for x in self.models] - for x in self.models: - x.train(True) - - def __enter__(self): - pass - - def __exit__(self, *args): - for x, mode in zip(self.models, self.prev_modes): - x.train(mode) - return False - - -def batch_to(dst, *args): - return [x.to(dst) if x is not None else None for x in args] - - -def get_flat_params_from(models): - if not hasattr(models, '__iter__'): - models = (models,) - params = [] - for model in models: - for param in model.parameters(): - params.append(param.data.view(-1)) - - flat_params = torch.cat(params) - return flat_params - - -def set_flat_params_to(model, flat_params): - prev_ind = 0 - for param in model.parameters(): - flat_size = int(np.prod(list(param.size()))) - param.data.copy_( - flat_params[prev_ind:prev_ind + flat_size].view(param.size())) - prev_ind += flat_size - - -def get_flat_grad_from(inputs, grad_grad=False): - grads = [] - for param in inputs: - if grad_grad: - grads.append(param.grad.grad.view(-1)) - else: - if param.grad is None: - grads.append(zeros(param.view(-1).shape)) - else: - grads.append(param.grad.view(-1)) - - flat_grad = torch.cat(grads) - return flat_grad - - -def compute_flat_grad(output, inputs, filter_input_ids=set(), retain_graph=False, create_graph=False): - if create_graph: - retain_graph = True - - inputs = list(inputs) - params = [] - for i, param in enumerate(inputs): - if i not in filter_input_ids: - params.append(param) - - grads = torch.autograd.grad(output, params, retain_graph=retain_graph, create_graph=create_graph) - - j = 0 - out_grads = [] - for i, param in enumerate(inputs): - if i in filter_input_ids: - out_grads.append(zeros(param.view(-1).shape)) - else: - out_grads.append(grads[j].view(-1)) - j += 1 - grads = torch.cat(out_grads) - - for param in params: - param.grad = None - return grads - - -def set_optimizer_lr(optimizer, lr): - for param_group in optimizer.param_groups: - param_group['lr'] = lr - - -def filter_state_dict(state_dict, filter_keys): - for key in list(state_dict.keys()): - for f_key in filter_keys: - if f_key in key: - del state_dict[key] - break - - -def get_scheduler(optimizer, policy, nepoch_fix=None, nepoch=None, decay_step=None): - if policy == 'lambda': - def lambda_rule(epoch): - lr_l = 1.0 - max(0, epoch - nepoch_fix) / float(nepoch - nepoch_fix + 1) - return lr_l - - scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) - elif policy == 'step': - scheduler = lr_scheduler.StepLR( - optimizer, step_size=decay_step, gamma=0.1) - elif policy == 'plateau': - scheduler = lr_scheduler.ReduceLROnPlateau( - optimizer, mode='min', factor=0.2, threshold=0.01, patience=5) - else: - return NotImplementedError('learning rate policy [%s] is not implemented', policy) - return scheduler - - -def init_weights(module): - if isinstance(module, nn.Conv2d): - nn.init.kaiming_normal_(module.weight.data, mode="fan_out") - elif isinstance(module, nn.BatchNorm2d): - nn.init.constant_(module.weight, 1.0) - nn.init.constant_(module.bias, 0.0) - elif isinstance(module, nn.Linear): - # print("weights ", module) - for name, param in module.named_parameters(): - if "bias" in name: - nn.init.constant_(param, 0.0) - elif "weight" in name: - nn.init.xavier_uniform_(param) - elif ( - isinstance(module, nn.LSTM) - or isinstance(module, nn.RNN) - or isinstance(module, nn.LSTMCell) - or isinstance(module, nn.RNNCell) - or isinstance(module, nn.GRU) - or isinstance(module, nn.GRUCell) - ): - # https://www.cse.iitd.ac.in/~mausam/courses/col772/spring2018/lectures/12-tricks.pdf - # • It can take a while for a RNN to learn to remember information - # • Initialize biases for LSTM’s forget gate to 1 to remember more by default. - # • Similarly, initialize biases for GRU’s reset gate to -1. - DIV = 3 if isinstance(module, nn.GRU) or isinstance(module, nn.GRUCell) else 4 - for name, param in module.named_parameters(): - if "bias" in name: - # print(name) - nn.init.constant_( - param, 0.0 - ) - if isinstance(module, nn.LSTMCell) \ - or isinstance(module, nn.LSTM): - n = param.size(0) - # LSTM: (W_ii|W_if|W_ig|W_io), W_if (forget gate) => bias 1 - start, end = n // DIV, n // 2 - param.data[start:end].fill_(1.) # to remember more by default - elif isinstance(module, nn.GRU) \ - or isinstance(module, nn.GRUCell): - # GRU: (W_ir|W_iz|W_in), W_ir (reset gate) => bias -1 - end = param.size(0) // DIV - param.data[:end].fill_(-1.) # to remember more by default - elif "weight" in name: - nn.init.xavier_normal_(param) - if isinstance(module, nn.LSTMCell) \ - or isinstance(module, nn.LSTM) \ - or isinstance(module, nn.GRU) \ - or isinstance(module, nn.GRUCell): - if 'weight_ih' in name: # input -> hidden weights - mul = param.shape[0] // DIV - for idx in range(DIV): - nn.init.xavier_uniform_(param[idx * mul:(idx + 1) * mul]) - elif 'weight_hh' in name: # hidden -> hidden weights (recurrent) - mul = param.shape[0] // DIV - for idx in range(DIV): - nn.init.orthogonal_(param[idx * mul:( - idx + 1) * mul]) # orthogonal initialization https://arxiv.org/pdf/1702.00071.pdf - else: - print(f"[WARNING] Module not initialized: {module}") diff --git a/personalised/code/framework/motion_diffusion/diffusion/utils/__init__.py b/personalised/code/framework/motion_diffusion/diffusion/utils/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/personalised/code/framework/motion_diffusion/diffusion/utils/geometry.py b/personalised/code/framework/motion_diffusion/diffusion/utils/geometry.py deleted file mode 100644 index 71b1c3caa86d6bf570b6b8f1edaa3549bade7bfc..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/utils/geometry.py +++ /dev/null @@ -1,473 +0,0 @@ -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import torch -import numpy as np -from torch.nn import functional as F - - -def matrix_of_angles(cos, sin, inv=False, dim=2): - assert dim in [2, 3] - sin = -sin if inv else sin - if dim == 2: - row1 = torch.stack((cos, -sin), axis=-1) - row2 = torch.stack((sin, cos), axis=-1) - return torch.stack((row1, row2), axis=-2) - elif dim == 3: - row1 = torch.stack((cos, -sin, 0 * cos), axis=-1) - row2 = torch.stack((sin, cos, 0 * cos), axis=-1) - row3 = torch.stack((0 * sin, 0 * cos, 1 + 0 * cos), axis=-1) - return torch.stack((row1, row2, row3), axis=-2) - - -def matrot2axisangle(matrots): - # This function is borrowed from https://github.com/davrempe/humor/utils/transforms.py - # axisang N x 3 - ''' - :param matrots: N*num_joints*9 - :return: N*num_joints*3 - ''' - import cv2 - batch_size = matrots.shape[0] - matrots = matrots.reshape([batch_size, -1, 9]) - out_axisangle = [] - for mIdx in range(matrots.shape[0]): - cur_axisangle = [] - for jIdx in range(matrots.shape[1]): - a = cv2.Rodrigues(matrots[mIdx, - jIdx:jIdx + 1, :].reshape(3, - 3))[0].reshape( - (1, 3)) - cur_axisangle.append(a) - - out_axisangle.append(np.array(cur_axisangle).reshape([1, -1, 3])) - return np.vstack(out_axisangle) - - -def axisangle2matrots(axisangle): - # This function is borrowed from https://github.com/davrempe/humor/utils/transforms.py - # axisang N x 3 - ''' - :param axisangle: N*num_joints*3 - :return: N*num_joints*9 - ''' - import cv2 - batch_size = axisangle.shape[0] - axisangle = axisangle.reshape([batch_size, -1, 3]) - out_matrot = [] - for mIdx in range(axisangle.shape[0]): - cur_axisangle = [] - for jIdx in range(axisangle.shape[1]): - a = cv2.Rodrigues(axisangle[mIdx, jIdx:jIdx + 1, :].reshape(1, - 3))[0] - cur_axisangle.append(a) - - out_matrot.append(np.array(cur_axisangle).reshape([1, -1, 9])) - return np.vstack(out_matrot) - - -def batch_rodrigues(axisang): - # This function is borrowed from https://github.com/MandyMo/pytorch_HMR/blob/master/src/util.py#L37 - # axisang N x 3 - axisang_norm = torch.norm(axisang + 1e-8, p=2, dim=1) - angle = torch.unsqueeze(axisang_norm, -1) - axisang_normalized = torch.div(axisang, angle) - angle = angle * 0.5 - v_cos = torch.cos(angle) - v_sin = torch.sin(angle) - - quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1) - rot_mat = quat2mat(quat) - rot_mat = rot_mat.view(rot_mat.shape[0], 9) - return rot_mat - - -def quat2mat(quat): - """ - This function is borrowed from https://github.com/MandyMo/pytorch_HMR/blob/master/src/util.py#L50 - - Convert quaternion coefficients to rotation matrix. - Args: - quat: size = [batch_size, 4] 4 <===>(w, x, y, z) - Returns: - Rotation matrix corresponding to the quaternion -- size = [batch_size, 3, 3] - """ - norm_quat = quat - norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True) - w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:, - 2], norm_quat[:, - 3] - - batch_size = quat.size(0) - - w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2) - wx, wy, wz = w * x, w * y, w * z - xy, xz, yz = x * y, x * z, y * z - - rotMat = torch.stack([ - w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz, 2 * wz + 2 * xy, - w2 - x2 + y2 - z2, 2 * yz - 2 * wx, 2 * xz - 2 * wy, 2 * wx + 2 * yz, - w2 - x2 - y2 + z2 - ], - dim=1).view(batch_size, 3, 3) - return rotMat - - -def rotation_matrix_to_angle_axis(rotation_matrix): - """ - This function is borrowed from https://github.com/kornia/kornia - - Convert 3x4 rotation matrix to Rodrigues vector - - Args: - rotation_matrix (Tensor): rotation matrix. - - Returns: - Tensor: Rodrigues vector transformation. - - Shape: - - Input: :math:`(N, 3, 4)` - - Output: :math:`(N, 3)` - - Example: - >>> input = torch.rand(2, 3, 4) # Nx4x4 - >>> output = tgm.rotation_matrix_to_angle_axis(input) # Nx3 - """ - if rotation_matrix.shape[1:] == (3, 3): - rot_mat = rotation_matrix.reshape(-1, 3, 3) - hom = torch.tensor([0, 0, 1], - dtype=torch.float32, - device=rotation_matrix.device).reshape( - 1, 3, 1).expand(rot_mat.shape[0], -1, -1) - rotation_matrix = torch.cat([rot_mat, hom], dim=-1) - - quaternion = rotation_matrix_to_quaternion(rotation_matrix) - aa = quaternion_to_angle_axis(quaternion) - aa[torch.isnan(aa)] = 0.0 - return aa - - -def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor: - """ - This function is borrowed from https://github.com/kornia/kornia - - Convert quaternion vector to angle axis of rotation. - - Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h - - Args: - quaternion (torch.Tensor): tensor with quaternions. - - Return: - torch.Tensor: tensor with angle axis of rotation. - - Shape: - - Input: :math:`(*, 4)` where `*` means, any number of dimensions - - Output: :math:`(*, 3)` - - Example: - >>> quaternion = torch.rand(2, 4) # Nx4 - >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion) # Nx3 - """ - if not torch.is_tensor(quaternion): - raise TypeError("Input type is not a torch.Tensor. Got {}".format( - type(quaternion))) - - if not quaternion.shape[-1] == 4: - raise ValueError( - "Input must be a tensor of shape Nx4 or 4. Got {}".format( - quaternion.shape)) - # unpack input and compute conversion - q1: torch.Tensor = quaternion[..., 1] - q2: torch.Tensor = quaternion[..., 2] - q3: torch.Tensor = quaternion[..., 3] - sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3 - - sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta) - cos_theta: torch.Tensor = quaternion[..., 0] - two_theta: torch.Tensor = 2.0 * torch.where( - cos_theta < 0.0, torch.atan2(-sin_theta, -cos_theta), - torch.atan2(sin_theta, cos_theta)) - - k_pos: torch.Tensor = two_theta / sin_theta - k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta) - k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg) - - angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3] - angle_axis[..., 0] += q1 * k - angle_axis[..., 1] += q2 * k - angle_axis[..., 2] += q3 * k - return angle_axis - - -def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6): - """ - This function is borrowed from https://github.com/kornia/kornia - - Convert 3x4 rotation matrix to 4d quaternion vector - - This algorithm is based on algorithm described in - https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201 - - Args: - rotation_matrix (Tensor): the rotation matrix to convert. - - Return: - Tensor: the rotation in quaternion - - Shape: - - Input: :math:`(N, 3, 4)` - - Output: :math:`(N, 4)` - - Example: - >>> input = torch.rand(4, 3, 4) # Nx3x4 - >>> output = tgm.rotation_matrix_to_quaternion(input) # Nx4 - """ - if not torch.is_tensor(rotation_matrix): - raise TypeError("Input type is not a torch.Tensor. Got {}".format( - type(rotation_matrix))) - - if len(rotation_matrix.shape) > 3: - raise ValueError( - "Input size must be a three dimensional tensor. Got {}".format( - rotation_matrix.shape)) - if not rotation_matrix.shape[-2:] == (3, 4): - raise ValueError( - "Input size must be a N x 3 x 4 tensor. Got {}".format( - rotation_matrix.shape)) - - rmat_t = torch.transpose(rotation_matrix, 1, 2) - - mask_d2 = rmat_t[:, 2, 2] < eps - - mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1] - mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1] - - t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2] - q0 = torch.stack([ - rmat_t[:, 1, 2] - rmat_t[:, 2, 1], t0, - rmat_t[:, 0, 1] + rmat_t[:, 1, 0], rmat_t[:, 2, 0] + rmat_t[:, 0, 2] - ], -1) - t0_rep = t0.repeat(4, 1).t() - - t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2] - q1 = torch.stack([ - rmat_t[:, 2, 0] - rmat_t[:, 0, 2], rmat_t[:, 0, 1] + rmat_t[:, 1, 0], - t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1] - ], -1) - t1_rep = t1.repeat(4, 1).t() - - t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2] - q2 = torch.stack([ - rmat_t[:, 0, 1] - rmat_t[:, 1, 0], rmat_t[:, 2, 0] + rmat_t[:, 0, 2], - rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2 - ], -1) - t2_rep = t2.repeat(4, 1).t() - - t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2] - q3 = torch.stack([ - t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1], - rmat_t[:, 2, 0] - rmat_t[:, 0, 2], rmat_t[:, 0, 1] - rmat_t[:, 1, 0] - ], -1) - t3_rep = t3.repeat(4, 1).t() - - mask_c0 = mask_d2 * mask_d0_d1 - mask_c1 = mask_d2 * ~mask_d0_d1 - mask_c2 = ~mask_d2 * mask_d0_nd1 - mask_c3 = ~mask_d2 * ~mask_d0_nd1 - mask_c0 = mask_c0.view(-1, 1).type_as(q0) - mask_c1 = mask_c1.view(-1, 1).type_as(q1) - mask_c2 = mask_c2.view(-1, 1).type_as(q2) - mask_c3 = mask_c3.view(-1, 1).type_as(q3) - - q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3 - q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 + # noqa - t2_rep * mask_c2 + t3_rep * mask_c3) # noqa - q *= 0.5 - return q - - -def estimate_translation_np(S, - joints_2d, - joints_conf, - focal_length=5000., - img_size=224.): - """ - This function is borrowed from https://github.com/nkolot/SPIN/utils/geometry.py - - Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d. - Input: - S: (25, 3) 3D joint locations - joints: (25, 3) 2D joint locations and confidence - Returns: - (3,) camera translation vector - """ - - num_joints = S.shape[0] - # focal length - f = np.array([focal_length, focal_length]) - # optical center - center = np.array([img_size / 2., img_size / 2.]) - - # transformations - Z = np.reshape(np.tile(S[:, 2], (2, 1)).T, -1) - XY = np.reshape(S[:, 0:2], -1) - O = np.tile(center, num_joints) - F = np.tile(f, num_joints) - weight2 = np.reshape(np.tile(np.sqrt(joints_conf), (2, 1)).T, -1) - - # least squares - Q = np.array([ - F * np.tile(np.array([1, 0]), num_joints), - F * np.tile(np.array([0, 1]), num_joints), - O - np.reshape(joints_2d, -1) - ]).T - c = (np.reshape(joints_2d, -1) - O) * Z - F * XY - - # weighted least squares - W = np.diagflat(weight2) - Q = np.dot(W, Q) - c = np.dot(W, c) - - # square matrix - A = np.dot(Q.T, Q) - b = np.dot(Q.T, c) - - # solution - trans = np.linalg.solve(A, b) - - return trans - - -def estimate_translation(S, joints_2d, focal_length=5000., img_size=224.): - """ - This function is borrowed from https://github.com/nkolot/SPIN/utils/geometry.py - - Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d. - Input: - S: (B, 49, 3) 3D joint locations - joints: (B, 49, 3) 2D joint locations and confidence - Returns: - (B, 3) camera translation vectors - """ - - device = S.device - # Use only joints 25:49 (GT joints) - S = S[:, 25:, :].cpu().numpy() - joints_2d = joints_2d[:, 25:, :].cpu().numpy() - joints_conf = joints_2d[:, :, -1] - joints_2d = joints_2d[:, :, :-1] - trans = np.zeros((S.shape[0], 3), dtype=np.float6432) - # Find the translation for each example in the batch - for i in range(S.shape[0]): - S_i = S[i] - joints_i = joints_2d[i] - conf_i = joints_conf[i] - trans[i] = estimate_translation_np(S_i, - joints_i, - conf_i, - focal_length=focal_length, - img_size=img_size) - return torch.from_numpy(trans).to(device) - - -def rot6d_to_rotmat_spin(x): - """Convert 6D rotation representation to 3x3 rotation matrix. - Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019 - Input: - (B,6) Batch of 6-D rotation representations - Output: - (B,3,3) Batch of corresponding rotation matrices - """ - x = x.view(-1, 3, 2) - a1 = x[:, :, 0] - a2 = x[:, :, 1] - b1 = F.normalize(a1) - b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1) - - # inp = a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1 - # denom = inp.pow(2).sum(dim=1).sqrt().unsqueeze(-1) + 1e-8 - # b2 = inp / denom - - b3 = torch.cross(b1, b2) - return torch.stack((b1, b2, b3), dim=-1) - - -def rot6d_to_rotmat(x): - x = x.view(-1, 3, 2) - - # Normalize the first vector - b1 = F.normalize(x[:, :, 0], dim=1, eps=1e-6) - - dot_prod = torch.sum(b1 * x[:, :, 1], dim=1, keepdim=True) - # Compute the second vector by finding the orthogonal complement to it - b2 = F.normalize(x[:, :, 1] - dot_prod * b1, dim=-1, eps=1e-6) - - # Finish building the basis by taking the cross product - b3 = torch.cross(b1, b2, dim=1) - rot_mats = torch.stack([b1, b2, b3], dim=-1) - - return rot_mats - - -# import mld.utils.rotation_conversions as rotation_conversions -# -# -# def rot6d(x_rotations, pose_rep): -# time, njoints, feats = x_rotations.shape -# -# # Compute rotations (convert only masked sequences output) -# if pose_rep == "rotvec": -# rotations = rotation_conversions.axis_angle_to_matrix(x_rotations) -# elif pose_rep == "rotmat": -# rotations = x_rotations.view(njoints, 3, 3) -# elif pose_rep == "rotquat": -# rotations = rotation_conversions.quaternion_to_matrix(x_rotations) -# elif pose_rep == "rot6d": -# rotations = rotation_conversions.rotation_6d_to_matrix(x_rotations) -# else: -# raise NotImplementedError("No geometry for this one.") -# -# rotations_6d = rotation_conversions.matrix_to_rotation_6d(rotations) -# return rotations_6d -# -# -# def rot6d_batch(x_rotations, pose_rep): -# nsamples, time, njoints, feats = x_rotations.shape -# -# # Compute rotations (convert only masked sequences output) -# if pose_rep == "rotvec": -# rotations = rotation_conversions.axis_angle_to_matrix(x_rotations) -# elif pose_rep == "rotmat": -# rotations = x_rotations.view(-1, njoints, 3, 3) -# elif pose_rep == "rotquat": -# rotations = rotation_conversions.quaternion_to_matrix(x_rotations) -# elif pose_rep == "rot6d": -# rotations = rotation_conversions.rotation_6d_to_matrix(x_rotations) -# else: -# raise NotImplementedError("No geometry for this one.") -# -# rotations_6d = rotation_conversions.matrix_to_rotation_6d(rotations) -# return rotations_6d -# -# -# def rot6d_to_rotvec_batch(pose): -# # nsamples, time, njoints, feats = rot6d.shape -# bs, nfeats = pose.shape -# rot6d = pose.reshape(bs, 24, 6) -# rotations = rotation_conversions.rotation_6d_to_matrix(rot6d) -# rotvec = rotation_conversions.matrix_to_axis_angle(rotations) -# return rotvec.reshape(bs, 24 * 3) diff --git a/personalised/code/framework/motion_diffusion/diffusion/utils/losses.py b/personalised/code/framework/motion_diffusion/diffusion/utils/losses.py deleted file mode 100644 index 3e757bc5c8189e52b574e9fd859cdc00cf0503ea..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/utils/losses.py +++ /dev/null @@ -1,69 +0,0 @@ -from __future__ import print_function - -import torch - - -def TemporalLoss(Y): - diff = Y[:, 1:, :] - Y[:, :-1, :] - t_loss = torch.mean(torch.norm(diff, dim=2, p=2) ** 2) - return t_loss - - -def L1Loss(prediction, target, k=1, reduction="min", **kwargs): - # prediction has shape of [batch_size, num_preds, features] - # target has shape of [batch_size, num_preds, features] - assert len(prediction.shape) == len(target.shape), "prediction and target must have the same shape" - assert len(prediction.shape) == 3, "Only works with predictions of shape [batch_size, num_preds, features]" - - # manual implementation of L1 loss - loss = (torch.abs(prediction - target)).mean(axis=-1) - - # reduce across multiple predictions - if reduction == "mean": - loss = torch.mean(loss) - elif reduction == "min": - loss = loss.min(axis=-1)[0].mean() - else: - raise NotImplementedError("reduction {} not implemented".format(reduction)) - return loss - - -def MSELoss(prediction, target, k=1, reduction="mean", **kwargs): - # prediction has shape of [batch_size, num_preds==k, features] - # target has shape of [batch_size, num_preds==k, features] - assert len(prediction.shape) == len(target.shape), "prediction and target must have the same shape" - assert len(prediction.shape) == 3, "Only works with predictions of shape [batch_size, num_preds, features]" - - # manual implementation of MSE loss - loss = ((prediction - target) ** 2).mean(axis=-1) # (batch_size, k) - - # reduce across multiple predictions - if reduction == "mean": - loss = torch.mean(loss) - elif reduction == "min": - loss = loss.min(axis=-1)[0].mean() - else: - raise NotImplementedError("reduction {} not implemented".format(reduction)) - return loss - - -class MSELoss_AE: - def __init__(self, w_mse=1, w_kld=1, w_coeff=1, **kwargs): - self.w_mse = w_mse - self.w_kld = w_kld - self.w_coeff = w_coeff - - def __call__(self, prediction, target, coefficients_3dmm, target_coefficients, mu, logvar): - batch_size = prediction.shape[0] - - prediction = prediction.reshape(prediction.shape[0], -1) - target = target.reshape(target.shape[0], -1) - coefficients_3dmm = coefficients_3dmm.reshape(coefficients_3dmm.shape[0], -1) - target_coefficients = target_coefficients.reshape(target_coefficients.shape[0], -1) - - MSE = ((prediction - target) ** 2).mean() - KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / batch_size - COEFF = ((coefficients_3dmm - target_coefficients) ** 2).mean() - - loss_r = self.w_mse * MSE + self.w_kld * KLD + self.w_coeff * COEFF - return {"loss": loss_r, "mse": MSE, "kld": KLD, "coeff": COEFF} diff --git a/personalised/code/framework/motion_diffusion/diffusion/utils/temos_utils.py b/personalised/code/framework/motion_diffusion/diffusion/utils/temos_utils.py deleted file mode 100644 index 717ed8f3ddbbecda431a5495d9f410db678899c9..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/utils/temos_utils.py +++ /dev/null @@ -1,131 +0,0 @@ -from typing import Dict, List -import numpy as np -import torch -from torch import Tensor -from framework.motion_diffusion.diffusion.utils import geometry - - -def lengths_to_mask(lengths: List[int], - device: torch.device, - max_len: int = None) -> Tensor: - lengths = torch.tensor(lengths, device=device) - max_len = max_len if max_len else max(lengths) - mask = torch.arange(max_len, device=device).expand( - len(lengths), max_len) < lengths.unsqueeze(1) - return mask - - -def detach_to_numpy(tensor): - return tensor.detach().cpu().numpy() - - -def remove_padding(tensors, lengths): - return [ - tensor[:tensor_length] - for tensor, tensor_length in zip(tensors, lengths) - ] - - -def nfeats_of(rottype): - if rottype in ["rotvec", "axisangle"]: - return 3 - elif rottype in ["rotquat", "quaternion"]: - return 4 - elif rottype in ["rot6d", "6drot", "rotation6d"]: - return 6 - elif rottype in ["rotmat"]: - return 9 - else: - return TypeError("This rotation type doesn't have features.") - - -def axis_angle_to(newtype, rotations): - if newtype in ["matrix"]: - rotations = geometry.axis_angle_to_matrix(rotations) - return rotations - elif newtype in ["rotmat"]: - rotations = geometry.axis_angle_to_matrix(rotations) - rotations = matrix_to("rotmat", rotations) - return rotations - elif newtype in ["rot6d", "6drot", "rotation6d"]: - rotations = geometry.axis_angle_to_matrix(rotations) - rotations = matrix_to("rot6d", rotations) - return rotations - elif newtype in ["rotquat", "quaternion"]: - rotations = geometry.axis_angle_to_quaternion(rotations) - return rotations - elif newtype in ["rotvec", "axisangle"]: - return rotations - else: - raise NotImplementedError - - -def matrix_to(newtype, rotations): - if newtype in ["matrix"]: - return rotations - if newtype in ["rotmat"]: - rotations = rotations.reshape((*rotations.shape[:-2], 9)) - return rotations - elif newtype in ["rot6d", "6drot", "rotation6d"]: - rotations = geometry.matrix_to_rotation_6d(rotations) - return rotations - elif newtype in ["rotquat", "quaternion"]: - rotations = geometry.matrix_to_quaternion(rotations) - return rotations - elif newtype in ["rotvec", "axisangle"]: - rotations = geometry.matrix_to_axis_angle(rotations) - return rotations - else: - raise NotImplementedError - - -def to_matrix(oldtype, rotations): - if oldtype in ["matrix"]: - return rotations - if oldtype in ["rotmat"]: - rotations = rotations.reshape((*rotations.shape[:-2], 3, 3)) - return rotations - elif oldtype in ["rot6d", "6drot", "rotation6d"]: - rotations = geometry.rotation_6d_to_matrix(rotations) - return rotations - elif oldtype in ["rotquat", "quaternion"]: - rotations = geometry.quaternion_to_matrix(rotations) - return rotations - elif oldtype in ["rotvec", "axisangle"]: - rotations = geometry.axis_angle_to_matrix(rotations) - return rotations - else: - raise NotImplementedError - - -# TODO: use a real subsampler.. -def subsample(num_frames, last_framerate, new_framerate): - step = int(last_framerate / new_framerate) - assert step >= 1 - frames = np.arange(0, num_frames, step) - return frames - - -# TODO: use a real upsampler.. -def upsample(motion, last_framerate, new_framerate): - step = int(new_framerate / last_framerate) - assert step >= 1 - - # Alpha blending => interpolation - alpha = np.linspace(0, 1, step + 1) - last = np.einsum("l,...->l...", 1 - alpha, motion[:-1]) - new = np.einsum("l,...->l...", alpha, motion[1:]) - - chuncks = (last + new)[:-1] - output = np.concatenate(chuncks.swapaxes(1, 0)) - # Don't forget the last one - output = np.concatenate((output, motion[[-1]])) - return output - - -if __name__ == "__main__": - motion = np.arange(105) - submotion = motion[subsample(len(motion), 100.0, 12.5)] - newmotion = upsample(submotion, 12.5, 100) - - print(newmotion) diff --git a/personalised/code/framework/motion_diffusion/diffusion/utils/util.py b/personalised/code/framework/motion_diffusion/diffusion/utils/util.py deleted file mode 100644 index 5fac735a73ef9958af99718ad6d8c7389d74cebe..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/utils/util.py +++ /dev/null @@ -1,69 +0,0 @@ -import math -import torch - - -def prob_mask_like(shape, prob, device): - if prob == 1: - return torch.ones(shape, device=device, dtype=torch.bool) - elif prob == 0: - return torch.zeros(shape, device=device, dtype=torch.bool) - else: - return torch.zeros(shape, device=device).float().uniform_(0, 1) < prob - - -# Temporal Bias, inspired by ALiBi: https://github.com/ofirpress/attention_with_linear_biases -def tgt_biased_mask(n_head, max_seq_len, period): - def get_slopes(n): - def get_slopes_power_of_2(n): - start = (2 ** (-2 ** -(math.log2(n) - 3))) - ratio = start - return [start * ratio ** i for i in range(n)] - - if math.log2(n).is_integer(): - return get_slopes_power_of_2(n) - else: - closest_power_of_2 = 2 ** math.floor(math.log2(n)) - return (get_slopes_power_of_2(closest_power_of_2) + - get_slopes(2 * closest_power_of_2)[0::2][:n - closest_power_of_2]) - - slopes = torch.Tensor(get_slopes(n_head)) - bias = torch.arange(start=0, end=max_seq_len, step=period).unsqueeze(1).repeat(1, period).view(-1) // (period) - bias = - torch.flip(bias, dims=[0]) - alibi = torch.zeros(max_seq_len, max_seq_len) - for i in range(max_seq_len): - alibi[i, :i + 1] = bias[-(i + 1):] - alibi = slopes.unsqueeze(1).unsqueeze(1) * alibi.unsqueeze(0) - mask = (torch.triu(torch.ones(max_seq_len, max_seq_len)) == 1).transpose(0, 1) - mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) - mask = mask.unsqueeze(0) + alibi - return mask - - -def memory_biased_mask(n_head, window_size, max_seq_len, period): - def get_slopes(n): - def get_slopes_power_of_2(n): - start = (2 ** (-2 ** -(math.log2(n) - 3))) - ratio = start - return [start * ratio ** i for i in range(n)] - - if math.log2(n).is_integer(): - return get_slopes_power_of_2(n) - else: - closest_power_of_2 = 2 ** math.floor(math.log2(n)) - return (get_slopes_power_of_2(closest_power_of_2) + - get_slopes(2 * closest_power_of_2)[0::2][:n - closest_power_of_2]) - - slopes = torch.Tensor(get_slopes(n_head)) - - bias = torch.arange(start=0, end=max_seq_len, step=period).unsqueeze(1).repeat(1, period).view(-1) // (period) - bias = - torch.flip(bias, dims=[0]) - alibi = torch.zeros(window_size, max_seq_len) - for i in range(window_size): - alibi[i, :max_seq_len - window_size + i + 1] = bias[window_size - i - 1:] - alibi = slopes.unsqueeze(1).unsqueeze(1) * alibi.unsqueeze(0) - mask = torch.triu(torch.ones(window_size, max_seq_len)) == 1 - mask = torch.flip(mask, dims=[0, 1]) - - mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) - mask = mask.unsqueeze(0) + alibi - return mask diff --git a/personalised/code/framework/motion_diffusion/diffusion/velocity_transform.py b/personalised/code/framework/motion_diffusion/diffusion/velocity_transform.py deleted file mode 100644 index ac422df10f4965a9d2971350b98c01393199b1d5..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/diffusion/velocity_transform.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Velocity (frame-delta) reparametrisation for the listener-emotion sequence. - -The PerFRDiff decoder diffusion natively models the listener emotion sequence -``e`` of shape ``(..., T, C)`` directly in value space. This module provides a -*bijective* map between ``e`` and a velocity representation ``d`` so the diffusion -can instead be trained/sampled in delta space (predicting frame-to-frame change), -analogous to epsilon-/v-prediction in standard diffusion: - - to_delta: d[..., 0, :] = e[..., 0, :] (absolute anchor) - d[..., t, :] = e[..., t, :] - e[..., t-1, :] (velocity) - from_delta: e = cumsum(d, time) - -Keeping the very first frame as an absolute anchor makes the transform invertible -(no separate "initial frame" predictor is needed) and bounds long-range drift, -while every other position is a pure velocity target. - -Time axis is assumed to be ``dim=-2`` and channels ``dim=-1`` (i.e. ``(..., T, C)``), -which matches every tensor the decoder matcher feeds through here -(``(B, T, C)`` and ``(B * num_preds, T, C)``). -""" - -import torch - - -def to_delta(x: torch.Tensor) -> torch.Tensor: - """Map an emotion sequence ``(..., T, C)`` to its velocity representation. - - The first time step is kept as an absolute value; the rest are first-order - differences. Autograd-safe (no in-place writes). - """ - first = x[..., :1, :] - diffs = x[..., 1:, :] - x[..., :-1, :] - return torch.cat([first, diffs], dim=-2) - - -def from_delta(d: torch.Tensor) -> torch.Tensor: - """Inverse of :func:`to_delta`: integrate the velocity over time. - - ``from_delta(to_delta(x)) == x`` (up to floating point). - """ - return torch.cumsum(d, dim=-2) diff --git a/personalised/code/framework/motion_diffusion/losses/losses.py b/personalised/code/framework/motion_diffusion/losses/losses.py deleted file mode 100644 index f892497b252a891974fec5145eb38777920f9fd0..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_diffusion/losses/losses.py +++ /dev/null @@ -1,62 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class SupConLoss(nn.Module): - def __init__(self, temperature=0.07, base_temperature=0.07): - super(SupConLoss, self).__init__() - self.temperature = temperature - self.base_temperature = base_temperature - - def forward(self, feature, label, device): - assert len(feature) == len(label), "The length of features and mask is inconsistent." - N = feature.shape[0] - - label = label.contiguous().view(-1, 1) # shape: (N, 1) - mask = torch.eq(label, label.T).float().to(device) # shape: (N, N) - - anchor_feature = feature - contrast_feature = feature - # compute logits - anchor_dot_contrast = torch.div( - torch.matmul(anchor_feature, contrast_feature.T), - self.temperature) # shape: (N, N) - - # for numerical stability - logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True) - logits = anchor_dot_contrast - logits_max.detach() - - logits_mask = torch.scatter( - torch.ones_like(mask), - 1, - torch.arange(N).view(-1, 1).to(device), - 0 - ) - - mask = mask * logits_mask - - # compute log_prob - exp_logits = torch.exp(logits) * logits_mask - log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True)) # shape: (N, N) - - mask_pos_pairs = mask.sum(1) - mask_pos_pairs = torch.where(mask_pos_pairs < 1e-6, 1, mask_pos_pairs) - mean_log_prob_pos = (mask * log_prob).sum(1) / mask_pos_pairs - - loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos - - return loss.mean() - - -# VAE loss for auto-encoding -class KLLoss(nn.Module): - def __init__(self): - super(KLLoss, self).__init__() - - def forward(self, q, p): - div = torch.distributions.kl_divergence(q, p) - return div.mean() - - def __repr__(self): - return "KLLoss()" \ No newline at end of file diff --git a/personalised/code/framework/motion_transvae/BasicBlock.py b/personalised/code/framework/motion_transvae/BasicBlock.py deleted file mode 100644 index 161e39943ae3449b349a53181f9a395443e29ac4..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_transvae/BasicBlock.py +++ /dev/null @@ -1,109 +0,0 @@ -import torch -import torch.nn as nn -import numpy as np -import math - - -class ConvBlock(nn.Module): - def __init__(self, in_planes=3, planes=128): - super(ConvBlock, self).__init__() - self.planes = planes - self.conv1 = nn.Conv3d(in_planes, planes // 4, kernel_size=(1, 7, 7), stride=(1, 2, 2), padding=(0, 3, 3), - bias=False) - self.in1 = nn.InstanceNorm3d(planes // 4) - self.maxpool = nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 0, 0)) - - self.conv2 = nn.Conv3d(planes // 4, planes, kernel_size=(1, 7, 7), stride=(1, 2, 2), padding=(0, 3, 3), - bias=False) - self.in2 = nn.InstanceNorm3d(planes) - self.relu = nn.ReLU(inplace=True) - - self.conv3 = nn.Conv3d(planes, planes, kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1), bias=False) - self.in3 = nn.InstanceNorm3d(planes) - self.relu = nn.ReLU(inplace=True) - - self.conv4 = nn.Conv3d(planes, planes, kernel_size=(1, 3, 3), stride=(1, 1, 1), padding=(0, 1, 1), bias=False) - self.in4 = nn.InstanceNorm3d(planes) - self.relu = nn.ReLU(inplace=True) - - self.conv5 = nn.Conv3d(planes, planes, kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1), bias=False) - self.in5 = nn.InstanceNorm3d(planes) - self.relu = nn.ReLU(inplace=True) - - def forward(self, x): - """ - input: - speaker_video_frames x: (batch_size, 3, seq_len, img_size, img_size) - - output: - speaker_temporal_tokens y: (batch_size, token_dim, seq_len) - - """ - - x = self.relu(self.in1(self.conv1(x))) - x = self.maxpool(x) - x = self.relu(self.in2(self.conv2(x))) - x = self.relu(self.in3(self.conv3(x))) - x = x + self.relu(self.in4(self.conv4(x))) - x = self.relu(self.in5(self.conv5(x))) - - y = x.mean(dim=-1).mean(dim=-1) - return y - - -class PositionalEncoding(nn.Module): - def __init__(self, d_model, dropout=0.1, max_len=50000, batch_first=True): - super().__init__() - self.batch_first = batch_first - - self.dropout = nn.Dropout(p=dropout) - - pe = torch.zeros(max_len, d_model) - position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) - div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model)) - pe[:, 0::2] = torch.sin(position * div_term) - pe[:, 1::2] = torch.cos(position * div_term) - pe = pe.unsqueeze(0).transpose(0, 1) - - self.register_buffer('pe', pe) - - def forward(self, x): - # not used in the final model - if self.batch_first: - x = x + self.pe.permute(1, 0, 2)[:, :x.shape[1], :] - else: - x = x + self.pe[:x.shape[0], :] - return self.dropout(x) - - -def lengths_to_mask(lengths, device): - lengths = torch.tensor(lengths, device=device) - max_len = max(lengths) - mask = torch.arange(max_len, device=device).expand(len(lengths), max_len) < lengths.unsqueeze(1) - return mask - - - -# Temporal Bias, inspired by ALiBi: https://github.com/ofirpress/attention_with_linear_biases -def init_biased_mask(n_head, max_seq_len, period): - def get_slopes(n): - def get_slopes_power_of_2(n): - start = (2**(-2**-(math.log2(n)-3))) - ratio = start - return [start*ratio**i for i in range(n)] - if math.log2(n).is_integer(): - return get_slopes_power_of_2(n) - else: - closest_power_of_2 = 2**math.floor(math.log2(n)) - return get_slopes_power_of_2(closest_power_of_2) + get_slopes(2*closest_power_of_2)[0::2][:n-closest_power_of_2] - slopes = torch.Tensor(get_slopes(n_head)) - bias = torch.arange(start=0, end=max_seq_len, step=period).unsqueeze(1).repeat(1,period).view(-1)//(period) - bias = - torch.flip(bias,dims=[0]) - alibi = torch.zeros(max_seq_len, max_seq_len) - for i in range(max_seq_len): - alibi[i, :i+1] = bias[-(i+1):] - alibi = slopes.unsqueeze(1).unsqueeze(1) * alibi.unsqueeze(0) - mask = (torch.triu(torch.ones(max_seq_len, max_seq_len)) == 1).transpose(0, 1) - mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) - mask = mask.unsqueeze(0) + alibi - return mask \ No newline at end of file diff --git a/personalised/code/framework/motion_transvae/TransformerVAE.py b/personalised/code/framework/motion_transvae/TransformerVAE.py deleted file mode 100644 index 60e56163f767f54ce72718624689602583fde1fa..0000000000000000000000000000000000000000 --- a/personalised/code/framework/motion_transvae/TransformerVAE.py +++ /dev/null @@ -1,606 +0,0 @@ -import torch -import torch.nn as nn -from torch import Tensor -from .BasicBlock import ConvBlock, PositionalEncoding, init_biased_mask - - -def lengths_to_mask(lengths, - device: torch.device, - max_len: int = None) -> Tensor: - lengths = torch.tensor(lengths, device=device) - max_len = max_len if max_len else max(lengths) - mask = torch.arange(max_len, device=device).expand( - len(lengths), max_len) < lengths.unsqueeze(1) - return mask - - -class EEGPredictionHead(nn.Module): - def __init__(self, input_dim, hidden_dim=256, output_dim=14, dropout=0.1): - super().__init__() - self.net = nn.Sequential( - nn.LayerNorm(input_dim), - nn.Linear(input_dim, hidden_dim), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(hidden_dim, output_dim), - ) - - def forward(self, x): - return self.net(x) - - def get_model_name(self): - return self.__class__.__name__ - - -class VideoEncoder(nn.Module): - def __init__(self, img_size=224, feature_dim=128, device='cpu'): - super(VideoEncoder, self).__init__() - - self.img_size = img_size - self.feature_dim = feature_dim - - self.Conv3D = ConvBlock(3, feature_dim) - self.fc = nn.Linear(feature_dim, feature_dim) - self.device = device - - def forward(self, video): - """ - input: - speaker_video_frames x: (batch_size, seq_len, 3, img_size, img_size) - - output: - speaker_temporal_tokens y: (batch_size, seq_len, token_dim) - - """ - - video_input = video.transpose(1, 2) # B C T H W - token_output = self.Conv3D(video_input).transpose(1, 2) - token_output = self.fc(token_output) # B T C - return token_output - - -class VAEModel(nn.Module): - def __init__(self, - in_channels: int, - latent_dim: int = 256, - **kwargs) -> None: - super(VAEModel, self).__init__() - - self.latent_dim = latent_dim - self.in_channels = in_channels - self.linear = nn.Linear(in_channels, latent_dim) - - seq_trans_encoder_layer = nn.TransformerEncoderLayer(d_model=latent_dim, - nhead=4, - dim_feedforward=latent_dim * 2, - dropout=0.1) - - self.seqTransEncoder = nn.TransformerEncoder(seq_trans_encoder_layer, num_layers=1) - self.mu_token = nn.Parameter(torch.randn(latent_dim)) - self.logvar_token = nn.Parameter(torch.randn(latent_dim)) - - def forward(self, input): - x = self.linear(input) # B T D - B, T, D = input.shape - - lengths = [len(item) for item in input] - - mu_token = torch.tile(self.mu_token, (B,)).reshape(B, 1, -1) - logvar_token = torch.tile(self.logvar_token, (B,)).reshape(B, 1, -1) - - x = torch.cat([mu_token, logvar_token, x], dim=1) - - x = x.permute(1, 0, 2) - - token_mask = torch.ones((B, 2), dtype=bool, device=input.get_device()) - mask = lengths_to_mask(lengths, input.get_device()) - - aug_mask = torch.cat((token_mask, mask), 1) - - x = self.seqTransEncoder(x, src_key_padding_mask=~aug_mask) - - mu = x[0] - logvar = x[1] - std = logvar.exp().pow(0.5) - dist = torch.distributions.Normal(mu, std) - motion_sample = self.sample_from_distribution(dist).to(input.get_device()) - return motion_sample, dist - - def sample_from_distribution(self, distribution): - return distribution.rsample() - - -class Decoder(nn.Module): - def __init__(self, output_3dmm_dim=58, output_emotion_dim=25, feature_dim=128, device='cpu', max_seq_len=751, - n_head=4, window_size=8, online=False): - super(Decoder, self).__init__() - - self.feature_dim = feature_dim - self.window_size = window_size - self.device = device - self.online = online - - self.vae_model = VAEModel(feature_dim, feature_dim) - - if self.online: - self.lstm = nn.LSTM(feature_dim, feature_dim, 1, batch_first=True) - self.linear_3d = nn.Linear(output_3dmm_dim, feature_dim) - self.linear_reaction = nn.Linear(feature_dim, feature_dim) - decoder_layer_3d = nn.TransformerDecoderLayer(d_model=feature_dim, nhead=4, dim_feedforward=2 * feature_dim, - batch_first=True) - self.listener_reaction_decoder_3d = nn.TransformerDecoder(decoder_layer_3d, num_layers=1) - - decoder_layer = nn.TransformerDecoderLayer(d_model=feature_dim, nhead=n_head, dim_feedforward=2 * feature_dim, - batch_first=True) - self.listener_reaction_decoder_1 = nn.TransformerDecoder(decoder_layer, num_layers=1) - self.listener_reaction_decoder_2 = nn.TransformerDecoder(decoder_layer, num_layers=1) - - self.biased_mask = init_biased_mask(n_head=n_head, max_seq_len=max_seq_len, period=max_seq_len) - - self.listener_reaction_3dmm_map_layer = nn.Linear(feature_dim, output_3dmm_dim) - self.listener_reaction_emotion_map_layer = nn.Sequential( - nn.Linear(feature_dim + output_3dmm_dim, feature_dim), - nn.Linear(feature_dim, output_emotion_dim) - ) - self.PE = PositionalEncoding(feature_dim) - - def forward(self, encoded_feature, past_reaction_3dmm=None, past_reaction_emotion=None): - B, TS = encoded_feature.shape[0], encoded_feature.shape[1] - if self.online: - TL = self.window_size - else: - TL = TS - motion_sample, dist = self.vae_model(encoded_feature) - time_queries = torch.zeros(B, TL, self.feature_dim, device=encoded_feature.get_device()) - time_queries = self.PE(time_queries) - _dev = encoded_feature.device - tgt_mask = self.biased_mask[:, :TL, :TL].clone().detach().to(device=_dev).repeat(B, 1, 1) - - listener_reaction = self.listener_reaction_decoder_1(tgt=time_queries, memory=motion_sample.unsqueeze(1), - tgt_mask=tgt_mask) - listener_reaction = self.listener_reaction_decoder_2(listener_reaction, listener_reaction, tgt_mask=tgt_mask) - - if self.online and (past_reaction_3dmm is not None): - past_reaction_3dmm = self.linear_3d(past_reaction_3dmm) - past_reaction_3dmm_last = past_reaction_3dmm[:, -1] - - tgt_mask = self.biased_mask[:, :(TL + past_reaction_3dmm.shape[1]), - :(TL + past_reaction_3dmm.shape[1])].detach().to(device=_dev).repeat(B, 1, 1) - all_3dmm = torch.cat((past_reaction_3dmm, self.linear_reaction(listener_reaction)), dim=1) - listener_3dmm_out = self.listener_reaction_decoder_3d(all_3dmm, all_3dmm, tgt_mask=tgt_mask) - frame_num = listener_3dmm_out.shape[1] - listener_3dmm_out = listener_3dmm_out[:, (frame_num - TL):] - - listener_3dmm_out, _ = self.lstm(listener_3dmm_out, - (past_reaction_3dmm_last.view(1, B, self.feature_dim).contiguous(), - past_reaction_3dmm_last.view(1, B, self.feature_dim).contiguous())) - - listener_3dmm_out = self.listener_reaction_3dmm_map_layer(listener_3dmm_out) - else: - listener_3dmm_out = self.listener_reaction_3dmm_map_layer(listener_reaction) - - listener_emotion_out = self.listener_reaction_emotion_map_layer( - torch.cat((listener_3dmm_out, listener_reaction), dim=-1)) - - return listener_3dmm_out, listener_emotion_out, dist - - def reset_window_size(self, window_size): - self.window_size = window_size - - -class SpeakerBehaviourEncoder(nn.Module): - def __init__(self, img_size=224, audio_dim=78, feature_dim=128, device='cpu'): - super(SpeakerBehaviourEncoder, self).__init__() - - self.img_size = img_size - self.audio_dim = audio_dim - self.feature_dim = feature_dim - self.device = device - - self.video_encoder = VideoEncoder(img_size=img_size, feature_dim=feature_dim, device=device) - self.audio_feature_map = nn.Linear(self.audio_dim, self.feature_dim) - self.fusion_layer = nn.Linear(self.feature_dim * 2, self.feature_dim) - - def forward(self, video, audio): - video_feature = self.video_encoder(video) - audio_feature = self.audio_feature_map(audio) - speaker_behaviour_feature = self.fusion_layer(torch.cat((video_feature, audio_feature), dim=-1)) - - return speaker_behaviour_feature - - -class TransformerVAE(nn.Module): - def __init__(self, img_size=224, audio_dim=78, output_3dmm_dim=58, output_emotion_dim=25, feature_dim=128, - seq_len=750, task='online', window_size=8, device='cuda', eeg_head=None, **kwargs): - super(TransformerVAE, self).__init__() - - self.img_size = img_size - self.feature_dim = feature_dim - self.output_3dmm_dim = output_3dmm_dim - self.output_emotion_dim = output_emotion_dim - self.seq_len = seq_len - self.online = True if task == 'online' else False - self.window_size = window_size - self.device = device - self.register_buffer('_dev_buf', torch.zeros(1)) - - self.speaker_behaviour_encoder = SpeakerBehaviourEncoder(img_size, audio_dim, feature_dim, device) - self.reaction_decoder = Decoder(output_3dmm_dim=output_3dmm_dim, output_emotion_dim=output_emotion_dim, - feature_dim=feature_dim, device=device, window_size=self.window_size, - online=self.online) - self.fusion = nn.Linear(feature_dim + self.output_3dmm_dim + self.output_emotion_dim, feature_dim) - self.eeg_head = None - self.eeg_head_pooling = "mean" - self.eeg_detach_prediction_emotion = True - self.eeg_use_speaker_audio = True - self.eeg_use_speaker_emotion = True - self.eeg_use_speaker_3dmm = True - self.eeg_use_prediction_emotion = True - self.eeg_speaker_audio_dim = 0 - self.eeg_speaker_emotion_dim = 0 - self.eeg_speaker_3dmm_dim = 0 - self.eeg_prediction_emotion_dim = 0 - if eeg_head is not None and eeg_head.get("enabled", False): - self.eeg_head_pooling = eeg_head.get("pooling", "mean") - self.eeg_detach_prediction_emotion = eeg_head.get("detach_prediction_emotion", True) - self.eeg_use_speaker_audio = eeg_head.get("use_speaker_audio", True) - self.eeg_use_speaker_emotion = eeg_head.get("use_speaker_emotion", True) - self.eeg_use_speaker_3dmm = eeg_head.get("use_speaker_3dmm", True) - self.eeg_use_prediction_emotion = eeg_head.get("use_prediction_emotion", True) - self.eeg_speaker_audio_dim = audio_dim if self.eeg_use_speaker_audio else 0 - self.eeg_speaker_emotion_dim = output_emotion_dim if self.eeg_use_speaker_emotion else 0 - self.eeg_speaker_3dmm_dim = output_3dmm_dim if self.eeg_use_speaker_3dmm else 0 - self.eeg_prediction_emotion_dim = output_emotion_dim if self.eeg_use_prediction_emotion else 0 - eeg_input_dim = ( - self.eeg_speaker_audio_dim - + self.eeg_speaker_emotion_dim - + self.eeg_speaker_3dmm_dim - + self.eeg_prediction_emotion_dim - ) - if eeg_input_dim <= 0: - raise ValueError("At least one EEG head input source must be enabled.") - self.eeg_head = EEGPredictionHead( - input_dim=eeg_head.get("input_dim", eeg_input_dim), - hidden_dim=eeg_head.get("hidden_dim", 256), - output_dim=eeg_head.get("output_dim", 14), - dropout=eeg_head.get("dropout", 0.5), - ) - - def freeze_except_eeg_head(self): - if self.eeg_head is None: - raise RuntimeError("Cannot train EEG head only because eeg_head is disabled.") - - for parameter in self.parameters(): - parameter.requires_grad = False - for parameter in self.eeg_head.parameters(): - parameter.requires_grad = True - - def set_eeg_head_train_mode(self): - if self.eeg_head is None: - raise RuntimeError("Cannot train EEG head only because eeg_head is disabled.") - - self.eval() - self.eeg_head.train() - - @staticmethod - def _lengths_to_list(motion_lengths, batch_size, fallback_length): - if motion_lengths is None: - return [fallback_length] * batch_size - if torch.is_tensor(motion_lengths): - values = motion_lengths.detach().cpu().flatten().tolist() - elif isinstance(motion_lengths, (list, tuple)): - values = [ - int(item.item() if torch.is_tensor(item) else item) - for item in motion_lengths - ] - else: - values = [int(motion_lengths)] - if len(values) == 1 and batch_size > 1: - values = values * batch_size - if len(values) < batch_size: - values = values + [fallback_length] * (batch_size - len(values)) - return [int(value) for value in values[:batch_size]] - - def _pool_one_eeg_feature(self, feature, expected_dim, device, dtype): - if feature is None or feature.numel() == 0: - return torch.zeros(expected_dim, device=device, dtype=dtype) - - feature = feature.to(device=device, dtype=dtype) - if feature.dim() == 1: - return feature - if feature.dim() == 2: - if self.eeg_head_pooling == "last": - return feature[-1] - if self.eeg_head_pooling == "mean": - return feature.mean(dim=0) - raise ValueError(f"Unknown EEG head pooling: {self.eeg_head_pooling}") - raise ValueError(f"Unsupported EEG feature shape: {feature.shape}") - - def _pool_eeg_sequence(self, feature, expected_dim, batch_size, device, dtype): - if expected_dim <= 0: - return None - if feature is None: - return torch.zeros(batch_size, expected_dim, device=device, dtype=dtype) - - if isinstance(feature, (list, tuple)): - pooled = [ - self._pool_one_eeg_feature(item, expected_dim, device, dtype) - for item in feature - ] - if len(pooled) == 0: - return torch.zeros(batch_size, expected_dim, device=device, dtype=dtype) - if len(pooled) < batch_size: - pooled.extend( - torch.zeros(expected_dim, device=device, dtype=dtype) - for _ in range(batch_size - len(pooled)) - ) - return torch.stack(pooled[:batch_size], dim=0) - - feature = feature.to(device=device, dtype=dtype) - if feature.numel() == 0: - return torch.zeros(batch_size, expected_dim, device=device, dtype=dtype) - if feature.dim() == 3: - if self.eeg_head_pooling == "last": - pooled = feature[:, -1] - elif self.eeg_head_pooling == "mean": - pooled = feature.mean(dim=1) - else: - raise ValueError(f"Unknown EEG head pooling: {self.eeg_head_pooling}") - elif feature.dim() == 2: - if feature.shape[0] == batch_size and feature.shape[-1] == expected_dim: - pooled = feature - else: - pooled = self._pool_one_eeg_feature(feature, expected_dim, device, dtype).unsqueeze(0) - elif feature.dim() == 1: - pooled = feature.unsqueeze(0) - else: - raise ValueError(f"Unsupported EEG condition shape: {feature.shape}") - - if pooled.shape[0] == 1 and batch_size > 1: - pooled = pooled.expand(batch_size, -1) - return pooled - - @staticmethod - def _last_eeg_value(sequence, mask, length, device, dtype): - if sequence is None or sequence.numel() == 0: - return None, None - - sequence = sequence.to(device=device, dtype=dtype) - if sequence.dim() == 1: - target = sequence - default_mask = torch.ones_like(target) - elif sequence.dim() == 2: - index = min(max(int(length) - 1, 0), sequence.shape[0] - 1) - target = sequence[index] - default_mask = torch.ones_like(target) - else: - raise ValueError(f"Unsupported EEG target shape: {sequence.shape}") - - if mask is None or mask.numel() == 0: - return target, default_mask - mask = mask.to(device=device, dtype=dtype) - if mask.dim() == 1: - return target, mask - if mask.dim() == 2: - index = min(max(int(length) - 1, 0), mask.shape[0] - 1) - return target, mask[index] - raise ValueError(f"Unsupported EEG mask shape: {mask.shape}") - - def _eeg_targets_from_sequences(self, listener_eeg_input, listener_eeg_mask, - motion_lengths, batch_size, device, dtype): - if listener_eeg_input is None: - return None, None - - fallback_length = self.seq_len - lengths = self._lengths_to_list(motion_lengths, batch_size, fallback_length) - targets = [] - masks = [] - - if isinstance(listener_eeg_input, (list, tuple)): - mask_items = listener_eeg_mask if isinstance(listener_eeg_mask, (list, tuple)) else [None] * len(listener_eeg_input) - for index, sequence in enumerate(listener_eeg_input[:batch_size]): - mask = mask_items[index] if index < len(mask_items) else None - target, target_mask = self._last_eeg_value(sequence, mask, lengths[index], device, dtype) - if target is None: - continue - targets.append(target) - masks.append(target_mask) - else: - listener_eeg_input = listener_eeg_input.to(device=device, dtype=dtype) - if listener_eeg_input.dim() == 3: - listener_eeg_mask = listener_eeg_mask.to(device=device, dtype=dtype) \ - if listener_eeg_mask is not None else None - for index in range(min(batch_size, listener_eeg_input.shape[0])): - mask = listener_eeg_mask[index] if listener_eeg_mask is not None else None - target, target_mask = self._last_eeg_value( - listener_eeg_input[index], mask, lengths[index], device, dtype) - targets.append(target) - masks.append(target_mask) - else: - target, target_mask = self._last_eeg_value( - listener_eeg_input, - listener_eeg_mask.to(device=device, dtype=dtype) if listener_eeg_mask is not None else None, - lengths[0], - device, - dtype, - ) - if target is not None: - targets.append(target) - masks.append(target_mask) - - if not targets: - return None, None - target_eeg = torch.stack(targets, dim=0) - target_eeg_mask = torch.stack(masks, dim=0) - if target_eeg.shape[0] == 1 and batch_size > 1: - target_eeg = target_eeg.expand(batch_size, -1) - target_eeg_mask = target_eeg_mask.expand(batch_size, -1) - return target_eeg, target_eeg_mask - - def _attach_eeg_outputs(self, listener_emotion_out, speaker_audio=None, speaker_emotion=None, - speaker_3dmm=None, listener_eeg_input=None, listener_eeg_mask=None, - motion_lengths=None): - if self.eeg_head is None: - return {} - if listener_emotion_out is None: - return {} - - batch_size = len(listener_emotion_out) if isinstance(listener_emotion_out, list) else listener_emotion_out.shape[0] - first_prediction = listener_emotion_out[0] if isinstance(listener_emotion_out, list) else listener_emotion_out - device = first_prediction.device - dtype = first_prediction.dtype - feature_list = [] - - speaker_audio_feature = self._pool_eeg_sequence( - speaker_audio, self.eeg_speaker_audio_dim, batch_size, device, dtype) - if speaker_audio_feature is not None: - feature_list.append(speaker_audio_feature) - - speaker_emotion_feature = self._pool_eeg_sequence( - speaker_emotion, self.eeg_speaker_emotion_dim, batch_size, device, dtype) - if speaker_emotion_feature is not None: - feature_list.append(speaker_emotion_feature) - - speaker_3dmm_feature = self._pool_eeg_sequence( - speaker_3dmm, self.eeg_speaker_3dmm_dim, batch_size, device, dtype) - if speaker_3dmm_feature is not None: - feature_list.append(speaker_3dmm_feature) - - prediction_emotion_feature = self._pool_eeg_sequence( - listener_emotion_out, self.eeg_prediction_emotion_dim, batch_size, device, dtype) - if prediction_emotion_feature is not None: - if self.eeg_detach_prediction_emotion: - prediction_emotion_feature = prediction_emotion_feature.detach() - feature_list.append(prediction_emotion_feature) - - prediction_eeg = self.eeg_head(torch.cat(feature_list, dim=-1)) - outputs = {"prediction_eeg": prediction_eeg} - - target_eeg, target_eeg_mask = self._eeg_targets_from_sequences( - listener_eeg_input, - listener_eeg_mask, - motion_lengths, - batch_size, - prediction_eeg.device, - prediction_eeg.dtype, - ) - if target_eeg is not None: - outputs["target_eeg"] = target_eeg - outputs["target_eeg_mask"] = target_eeg_mask - return outputs - - def forward(self, speaker_video=None, speaker_audio=None, **kwargs): - speaker_emotion = kwargs.get("speaker_emotion", None) - speaker_3dmm = kwargs.get("speaker_3dmm", None) - listener_eeg_input = kwargs.get("listener_eeg_input", None) - listener_eeg_mask = kwargs.get("listener_eeg_mask", None) - return_eeg_outputs = kwargs.get("return_eeg_outputs", False) - return_distribution = kwargs.get("return_distribution", True) - distribution = [] if return_distribution else None - if self.online: - _dev = self._dev_buf.device - speaker_video = torch.stack(speaker_video, dim=0).to(device=_dev) - speaker_audio = torch.stack(speaker_audio, dim=0).to(device=_dev) - motion_lengths = ( - torch.as_tensor( - kwargs.get('motion_lengths', [self.seq_len] * len(speaker_video)), device=_dev, - ).clamp(max=self.seq_len) - ) - - frame_num = speaker_video.shape[1] - period = frame_num // self.window_size - # num_windows = motion_lengths // self.window_size - # motion_lengths Tensor([58, 720, 625, 750, ...]) ==> num_windows Tensor([7, 90, 78, 93, ...]) - - reaction_3dmm = torch.zeros((speaker_video.size(0), self.window_size, self.output_3dmm_dim), - device=_dev) - reaction_emotion = torch.zeros((speaker_video.size(0), self.window_size, self.output_emotion_dim), - device=_dev) - - for i in range(0, period): - # mask = (~(num_windows < i)).view(-1, 1).repeat(1, self.window_size) - # mask = mask.unsqueeze(-1).to(device=speaker_video) - speaker_video_, speaker_audio_ = (speaker_video[:, :(i + 1) * self.window_size], - speaker_audio[:, :(i + 1) * self.window_size]) - encoded_feature = self.speaker_behaviour_encoder(speaker_video_, speaker_audio_) - - # modality fusion - encoded_feature = self.fusion( - torch.cat((encoded_feature, reaction_3dmm, reaction_emotion), dim=-1)) - - if i != 0: - past_reaction_3dmm, past_reaction_emotion = (reaction_3dmm[:, :i * self.window_size], - reaction_emotion[:, :i * self.window_size]) - current_reaction_3dmm, current_reaction_emotion = (reaction_3dmm[:, i * self.window_size:], - reaction_emotion[:, i * self.window_size:]) - listener_3dmm_out, listener_emotion_out, dist = self.reaction_decoder(encoded_feature, - past_reaction_3dmm) - - reaction_3dmm = torch.cat( - (past_reaction_3dmm, listener_3dmm_out, current_reaction_3dmm), dim=1) - reaction_emotion = torch.cat( - (past_reaction_emotion, listener_emotion_out, current_reaction_emotion), dim=1) - - else: - listener_3dmm_out, listener_emotion_out, dist = self.reaction_decoder(encoded_feature) - reaction_3dmm = torch.cat((listener_3dmm_out, reaction_3dmm), dim=1) - reaction_emotion = torch.cat((listener_emotion_out, reaction_emotion), dim=1) - - if return_distribution: - distribution.append(dist) - - listener_3dmm_out, listener_emotion_out = reaction_3dmm[:, :frame_num], reaction_emotion[:, :frame_num] - seq_mask = lengths_to_mask(lengths=motion_lengths, - device=speaker_video.device, - max_len=frame_num).unsqueeze(-1).float() - listener_3dmm_out = listener_3dmm_out * seq_mask - listener_emotion_out = listener_emotion_out * seq_mask - - listener_3dmm_out, listener_emotion_out = \ - list(listener_3dmm_out), list(listener_emotion_out) - eeg_outputs = self._attach_eeg_outputs( - listener_emotion_out, - speaker_audio=speaker_audio, - speaker_emotion=speaker_emotion, - speaker_3dmm=speaker_3dmm, - listener_eeg_input=listener_eeg_input, - listener_eeg_mask=listener_eeg_mask, - motion_lengths=motion_lengths, - ) - if return_eeg_outputs: - return listener_3dmm_out, listener_emotion_out, distribution, eeg_outputs - return listener_3dmm_out, listener_emotion_out, distribution - - else: - _dev = self._dev_buf.device - listener_3dmm_outs = [] - listener_emotion_outs = [] - for speaker_video_, speaker_audio_ in zip(speaker_video, speaker_audio): # motion_lengths - speaker_video_, speaker_audio_ = \ - speaker_video_.unsqueeze(0).to(_dev), speaker_audio_.unsqueeze(0).to(_dev) - encoded_feature = self.speaker_behaviour_encoder(speaker_video_, speaker_audio_) - listener_3dmm_out, listener_emotion_out, dist = self.reaction_decoder(encoded_feature) - listener_3dmm_outs.append(listener_3dmm_out.squeeze(0)) - listener_emotion_outs.append(listener_emotion_out.squeeze(0)) - if return_distribution: - distribution.append(dist) - - eeg_outputs = self._attach_eeg_outputs( - listener_emotion_outs, - speaker_audio=speaker_audio, - speaker_emotion=speaker_emotion, - speaker_3dmm=speaker_3dmm, - listener_eeg_input=listener_eeg_input, - listener_eeg_mask=listener_eeg_mask, - motion_lengths=kwargs.get('motion_lengths', None), - ) - if return_eeg_outputs: - return listener_3dmm_outs, listener_emotion_outs, distribution, eeg_outputs - return listener_3dmm_outs, listener_emotion_outs, distribution - - def reset_window_size(self, window_size): - self.window_size = window_size - self.reaction_decoder.reset_window_size(window_size) - - def get_model_name(self): - return 'TransformerVAE' diff --git a/personalised/code/framework/motion_transvae/__init__.py b/personalised/code/framework/motion_transvae/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/personalised/code/framework/perfrdiff_rewrite_weight/__init__.py b/personalised/code/framework/perfrdiff_rewrite_weight/__init__.py deleted file mode 100644 index 2123b6ca52aa758e224c1db587d25c20c9f8219c..0000000000000000000000000000000000000000 --- a/personalised/code/framework/perfrdiff_rewrite_weight/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .modifier.network import MainNetUnified, ModifierNetwork -from .person_specific.PersonSpecificEncoder import Transformer diff --git a/personalised/code/framework/perfrdiff_rewrite_weight/losses.py b/personalised/code/framework/perfrdiff_rewrite_weight/losses.py deleted file mode 100644 index 8c16d6c2c90daa0aa8e3a6ff0ca768a8728fff05..0000000000000000000000000000000000000000 --- a/personalised/code/framework/perfrdiff_rewrite_weight/losses.py +++ /dev/null @@ -1,91 +0,0 @@ -import torch - - -def temporal_loss(y): - diff = y[:, 1:, :] - y[:, :-1, :] - return torch.mean(torch.norm(diff, dim=2, p=2) ** 2) - - -def l1_loss(prediction, target, reduction="min", **kwargs): - assert prediction.shape == target.shape - assert prediction.dim() == 3 - loss = torch.abs(prediction - target).mean(dim=-1) - if reduction == "mean": - return loss.mean() - if reduction == "min": - return loss.min(dim=-1)[0].mean() - raise NotImplementedError(f"Unsupported reduction: {reduction}") - - -def mse_loss(prediction, target, reduction="mean", **kwargs): - assert prediction.shape == target.shape - assert prediction.dim() == 3 - loss = ((prediction - target) ** 2).mean(dim=-1) - if reduction == "mean": - return loss.mean() - if reduction == "min": - return loss.min(dim=-1)[0].mean() - raise NotImplementedError(f"Unsupported reduction: {reduction}") - - -def masked_mse_loss(prediction, target, mask): - mask = mask.to(dtype=prediction.dtype) - loss = ((prediction - target) ** 2) * mask - return loss.sum() / mask.sum().clamp_min(1.0) - - -def DiffusionLoss( - output_prior, - output_decoder, - losses_type=("MSELoss", "MSELoss"), - losses_multipliers=(0.0, 1.0), - losses_decoded=(False, True), - k=1, - temporal_loss_w=0.0, - eeg_loss_weight=1.0, - **kwargs): - encoded_prediction = output_prior["encoded_prediction"] - encoded_target = output_prior["encoded_target"] - if encoded_prediction.dim() == 4: - encoded_prediction = encoded_prediction.squeeze(-2) - if encoded_target.dim() == 4: - encoded_target = encoded_target.squeeze(-2) - prediction_emotion = output_decoder["prediction_emotion"] - target_emotion = output_decoder["target_emotion"] - - if prediction_emotion.dim() == 4: - _, _, window_size, emotion_dim = prediction_emotion.shape - else: - _, window_size, emotion_dim = prediction_emotion.shape - - losses_dict = {"loss": prediction_emotion.new_tensor(0.0)} - losses_dict["loss_eeg"] = prediction_emotion.new_tensor(0.0) - losses_dict["eeg_valid_ratio"] = prediction_emotion.new_tensor(0.0) - losses_dict["temporal_loss"] = temporal_loss(prediction_emotion.reshape(-1, window_size, emotion_dim)) - losses_dict["loss"] = losses_dict["loss"] + losses_dict["temporal_loss"] * temporal_loss_w - - prediction_emotion = prediction_emotion.reshape(-1, k, window_size * emotion_dim) - target_emotion = target_emotion.reshape(-1, k, window_size * emotion_dim) - - loss_fns = { - "MSELoss": mse_loss, - "L1Loss": l1_loss, - } - for loss_name, weight, decoded in zip(losses_type, losses_multipliers, losses_decoded): - key = "decoded" if decoded else "encoded" - loss_fn = loss_fns[loss_name] - if decoded: - losses_dict[key] = loss_fn(prediction_emotion, target_emotion, k=k) - else: - losses_dict[key] = loss_fn(encoded_prediction, encoded_target, k=k) - losses_dict["loss"] = losses_dict["loss"] + losses_dict[key] * weight - - if "prediction_eeg" in output_decoder and "target_eeg" in output_decoder: - prediction_eeg = output_decoder["prediction_eeg"] - target_eeg = output_decoder["target_eeg"] - target_eeg_mask = output_decoder.get("target_eeg_mask", torch.ones_like(target_eeg)) - losses_dict["loss_eeg"] = masked_mse_loss(prediction_eeg, target_eeg, target_eeg_mask) - losses_dict["eeg_valid_ratio"] = target_eeg_mask.float().mean() - losses_dict["loss"] = losses_dict["loss"] + eeg_loss_weight * losses_dict["loss_eeg"] - - return losses_dict diff --git a/personalised/code/framework/perfrdiff_rewrite_weight/modifier/__init__.py b/personalised/code/framework/perfrdiff_rewrite_weight/modifier/__init__.py deleted file mode 100644 index 960da7aed698d14d6a381265c61127e3f2ed1298..0000000000000000000000000000000000000000 --- a/personalised/code/framework/perfrdiff_rewrite_weight/modifier/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .network import MainNetUnified, ModifierNetwork diff --git a/personalised/code/framework/perfrdiff_rewrite_weight/modifier/network.py b/personalised/code/framework/perfrdiff_rewrite_weight/modifier/network.py deleted file mode 100644 index 84cae73b2373b0d208bceb1835d94dc3f2ea3ab1..0000000000000000000000000000000000000000 --- a/personalised/code/framework/perfrdiff_rewrite_weight/modifier/network.py +++ /dev/null @@ -1,426 +0,0 @@ -from collections import OrderedDict - -import hydra -import os -import torch -import torch.nn as nn -import torch.nn.functional as F - -from framework.perfrdiff_rewrite_weight.person_specific.PersonSpecificEncoder import Transformer -from framework.utils.util import from_pretrained_checkpoint - - -def compute_regular_loss(weights): - loss = weights[0].new_tensor(0.0) - for weight in weights: - loss = loss + torch.norm(weight.reshape(-1), 2) - return loss - - -class ModifierNetwork(nn.Module): - def __init__(self, input_dim=512, latent_dim=1024, output_dim=None, num_shared_layers=1): - super().__init__() - output_dim = output_dim or [] - self.shared_layers = nn.ModuleList( - [ - nn.Linear(input_dim, latent_dim) if idx == 0 else nn.Linear(latent_dim, latent_dim) - for idx in range(num_shared_layers) - ] - ) - self.output_dim = output_dim - self.branches = nn.ModuleList( - [nn.Linear(latent_dim, int(torch.prod(shape).item())) for shape in output_dim] - ) - - def forward(self, x): - for layer in self.shared_layers: - x = torch.relu(layer(x)) - return [ - branch(x).view([int(dim) for dim in self.output_dim[idx]]) - for idx, branch in enumerate(self.branches) - ] - - def get_model_name(self): - return self.__class__.__name__ - - -class PersonalityEncoder(nn.Module): - def __init__(self, input_dim=5, hidden_dim=128, output_dim=512, dropout=0.1): - super().__init__() - self.net = nn.Sequential( - nn.LayerNorm(input_dim), - nn.Linear(input_dim, hidden_dim), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(hidden_dim, output_dim), - ) - - def forward(self, x): - return F.normalize(self.net(x), dim=-1) - - -class PersonalityFusion(nn.Module): - def __init__(self, embed_dim=512, hidden_dim=512, dropout=0.1): - super().__init__() - self.net = nn.Sequential( - nn.LayerNorm(embed_dim * 2), - nn.Linear(embed_dim * 2, hidden_dim), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(hidden_dim, embed_dim), - ) - - def forward(self, history_embedding, personality_embedding): - fused = self.net(torch.cat((history_embedding, personality_embedding), dim=-1)) - return F.normalize(fused, dim=-1) - - -class PersonCoarseConditioner(nn.Module): - """P4: map a person embedding -> FiLM (gamma, beta) over the causal coarse-GRU - hidden state, personalizing the 8-class expression *plan* directly (our - contribution, distinct from weight-editing the attention layers). The final - projection is zero-initialised so the generic plan is recovered exactly at the - start of training (gamma=0, beta=0).""" - - def __init__(self, embed_dim=512, coarse_hidden=256, hidden_dim=256, dropout=0.1): - super().__init__() - self.net = nn.Sequential( - nn.LayerNorm(embed_dim), - nn.Linear(embed_dim, hidden_dim), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(hidden_dim, 2 * coarse_hidden), - ) - nn.init.zeros_(self.net[-1].weight) - nn.init.zeros_(self.net[-1].bias) - - def forward(self, person_embedding): - gamma, beta = self.net(person_embedding).chunk(2, dim=-1) # each (B, H) - return gamma, beta - - -class MainNetUnified(nn.Module): - def __init__(self, cfg, main_net, device): - super().__init__() - self.main_net = main_net - self.modified_layers = list(cfg.main_model.args.modified_layers) - self.hypernet_predict = cfg.main_model.args.get("predict", "shift") - self.crossattn_modify = cfg.main_model.args.get("modify", "all") - self.regularization = cfg.main_model.args.get("regularization", False) - self.regular_w = cfg.main_model.args.get("regular_w", 0.0) - self.embed_dim = cfg.main_model.args.get("embed_dim", 512) - self.personal_condition_mode = cfg.main_model.args.get("personal_condition_mode", "3dmm_personality") - if self.personal_condition_mode not in {"3dmm_personality", "personality_only", "3dmm_only"}: - raise ValueError(f"Unsupported personal_condition_mode: {self.personal_condition_mode}") - - for parameter in self.main_net.parameters(): - parameter.requires_grad = False - - modules = OrderedDict(self.main_net.named_modules()) - missing = [name for name in self.modified_layers if name not in modules] - if missing: - raise ValueError(f"Cannot find modified layers in diffusion model: {missing}") - self.hooked_modules = OrderedDict((name, modules[name]) for name in self.modified_layers) - - weight_shapes = self.main_net.obtain_shapes(self.modified_layers) - missing_shapes = [name for name in self.modified_layers if name not in weight_shapes] - if missing_shapes: - raise ValueError(f"Cannot infer weights for modified layers: {missing_shapes}") - - self.weight_shapes = [] - for layer_name in self.modified_layers: - shape = weight_shapes[layer_name] - if "multihead_attn" in layer_name and self.crossattn_modify == "kv": - original_dim = shape[0] - shape = torch.tensor([2 * torch.div(original_dim, 3, rounding_mode="trunc"), shape[1]]) - self.weight_shapes.append(shape) - - self.hypernet = ModifierNetwork( - input_dim=cfg.main_model.args.get("input_dim", 512), - latent_dim=cfg.main_model.args.get("latent_dim", 1024), - output_dim=self.weight_shapes, - num_shared_layers=cfg.main_model.args.get("num_shared_layers", 1), - ) - - self.personality_encoder = None - self.personality_fusion = None - if self.personal_condition_mode in {"3dmm_personality", "personality_only"}: - personality_input_dim = cfg.main_model.args.get("personality_input_dim", 5) - personality_hidden_dim = cfg.main_model.args.get("personality_hidden_dim", 128) - personality_dropout = cfg.main_model.args.get("personality_dropout", 0.1) - self.personality_encoder = PersonalityEncoder( - input_dim=personality_input_dim, - hidden_dim=personality_hidden_dim, - output_dim=self.embed_dim, - dropout=personality_dropout, - ) - if self.personal_condition_mode == "3dmm_personality": - fusion_hidden_dim = cfg.main_model.args.get("personality_fusion_hidden_dim", self.embed_dim) - self.personality_fusion = PersonalityFusion( - embed_dim=self.embed_dim, - hidden_dim=fusion_hidden_dim, - dropout=personality_dropout, - ) - - self.person_encoder = None - if self.personal_condition_mode in {"3dmm_personality", "3dmm_only"}: - person_cfg = cfg.person_specific - self.person_encoder = Transformer(device, **person_cfg.args) - checkpoint_path = hydra.utils.to_absolute_path(person_cfg.checkpoint_path) - if not os.path.isfile(checkpoint_path): - raise FileNotFoundError( - f"Missing person-specific encoder checkpoint: {checkpoint_path}. " - "Please place it under pretrained_models/person_specific/ or override " - "trainer.person_specific.checkpoint_path." - ) - from_pretrained_checkpoint(checkpoint_path, self.person_encoder, device) - self.person_encoder.eval() - for parameter in self.person_encoder.parameters(): - parameter.requires_grad = False - - # (P4) optional personalization of OUR coarse expression plan: a small - # trainable conditioner turns the person embedding into a FiLM over the - # causal coarse-GRU hidden state inside the denoiser. Orthogonal to (and - # composable with) the weight-editing hypernet above. - self.person_coarse = None - self._coarse_denoiser = None - if cfg.main_model.args.get("personalize_coarse", False): - denoiser = modules.get("diffusion_decoder.model", None) - if denoiser is None or not getattr(denoiser, "use_coarse", False): - raise ValueError( - "personalize_coarse=True requires a causal coarse denoiser at " - "diffusion_decoder.model (use CausalLatentMatcher with use_coarse=true)." - ) - coarse_hidden = denoiser.coarse_gru.hidden_size - self.person_coarse = PersonCoarseConditioner( - embed_dim=self.embed_dim, - coarse_hidden=coarse_hidden, - hidden_dim=cfg.main_model.args.get("person_coarse_hidden", coarse_hidden), - dropout=cfg.main_model.args.get("person_coarse_dropout", 0.1), - ) - self._coarse_denoiser = denoiser - - self._initialize_editable_weights(device) - - def _initialize_editable_weights(self, device): - original_weights = [] - weight_kinds = [] - for name, module in self.hooked_modules.items(): - if hasattr(module, "weight"): - original_weights.append(module.weight.detach()) - weight_kinds.append("weight") - del module._parameters["weight"] - elif hasattr(module, "in_proj_weight"): - original_weights.append(module.in_proj_weight.detach()) - weight_kinds.append("in_proj_weight") - del module._parameters["in_proj_weight"] - else: - raise ValueError(f"Layer has no editable weight: {name}") - self.original_weights = original_weights - self.weight_kinds = weight_kinds - - self.tensor_0 = torch.zeros(size=(self.embed_dim, self.embed_dim), device=device) - self.tensor_1 = torch.tensor(1.0, device=device) - - def eeg_head(self): - return getattr(self.main_net, "eeg_head", None) - - def has_eeg_head(self): - return self.eeg_head() is not None - - def set_eeg_head_requires_grad(self, requires_grad=True): - if not self.has_eeg_head(): - raise RuntimeError("Cannot train/evaluate EEG because main_net.eeg_head is disabled.") - for parameter in self.eeg_head().parameters(): - parameter.requires_grad = requires_grad - - def freeze_except_eeg_head(self): - if not self.has_eeg_head(): - raise RuntimeError("Cannot train EEG head only because main_net.eeg_head is disabled.") - for parameter in self.parameters(): - parameter.requires_grad = False - self.set_eeg_head_requires_grad(True) - - def set_eeg_head_train_mode(self): - if not self.has_eeg_head(): - raise RuntimeError("Cannot train EEG head only because main_net.eeg_head is disabled.") - self.eval() - self.eeg_head().train() - - def modifier_parameters(self, include_eeg_head=False): - for parameter in self.hypernet.parameters(): - yield parameter - if self.personality_encoder is not None: - for parameter in self.personality_encoder.parameters(): - yield parameter - if self.personality_fusion is not None: - for parameter in self.personality_fusion.parameters(): - yield parameter - if self.person_coarse is not None: - for parameter in self.person_coarse.parameters(): - yield parameter - if include_eeg_head: - self.set_eeg_head_requires_grad(True) - for parameter in self.eeg_head().parameters(): - yield parameter - - def modifier_state_dict(self, include_eeg_head=False): - state_dict = { - f"hypernet.{name}": value - for name, value in self.hypernet.state_dict().items() - } - if self.personality_encoder is not None: - state_dict.update( - { - f"personality_encoder.{name}": value - for name, value in self.personality_encoder.state_dict().items() - } - ) - if self.personality_fusion is not None: - state_dict.update( - { - f"personality_fusion.{name}": value - for name, value in self.personality_fusion.state_dict().items() - } - ) - if self.person_coarse is not None: - state_dict.update( - { - f"person_coarse.{name}": value - for name, value in self.person_coarse.state_dict().items() - } - ) - if include_eeg_head and self.has_eeg_head(): - state_dict.update( - { - f"eeg_head.{name}": value - for name, value in self.eeg_head().state_dict().items() - } - ) - return state_dict - - def load_modifier_state_dict(self, state_dict): - if any(key.startswith("hypernet.") for key in state_dict): - hypernet_state = { - key[len("hypernet."):]: value - for key, value in state_dict.items() - if key.startswith("hypernet.") - } - self.hypernet.load_state_dict(hypernet_state) - - personality_state = { - key[len("personality_encoder."):]: value - for key, value in state_dict.items() - if key.startswith("personality_encoder.") - } - if personality_state: - if self.personality_encoder is None: - raise ValueError( - "Checkpoint contains personality_encoder weights, " - f"but current personal_condition_mode is {self.personal_condition_mode}." - ) - self.personality_encoder.load_state_dict(personality_state) - - if self.personality_fusion is not None: - fusion_state = { - key[len("personality_fusion."):]: value - for key, value in state_dict.items() - if key.startswith("personality_fusion.") - } - if fusion_state: - self.personality_fusion.load_state_dict(fusion_state) - person_coarse_state = { - key[len("person_coarse."):]: value - for key, value in state_dict.items() - if key.startswith("person_coarse.") - } - if person_coarse_state: - if self.person_coarse is None: - raise ValueError( - "Checkpoint contains person_coarse weights, but " - "personalize_coarse is disabled in the current config." - ) - self.person_coarse.load_state_dict(person_coarse_state) - eeg_state = { - key[len("eeg_head."):]: value - for key, value in state_dict.items() - if key.startswith("eeg_head.") - } - if eeg_state: - if not self.has_eeg_head(): - raise ValueError("Checkpoint contains eeg_head weights, but main_net.eeg_head is disabled.") - self.eeg_head().load_state_dict(eeg_state) - return - - self.hypernet.load_state_dict(state_dict) - - def encode_person_condition(self, p=None, personality=None): - if self.personal_condition_mode == "3dmm_only": - if p is None or p.numel() == 0: - raise ValueError("3dmm_only mode requires a non-empty personal 3DMM history.") - with torch.no_grad(): - _, person_embedding = self.person_encoder(p) - else: - if personality is None or personality.numel() == 0: - raise ValueError("listener personality is required for perfrdiff personal conditioning.") - if personality.dim() == 1: - personality = personality.unsqueeze(0) - personality_embedding = self.personality_encoder(personality.float()) - if self.personal_condition_mode == "personality_only": - person_embedding = personality_embedding - else: - if p is None or p.numel() == 0: - raise ValueError("3dmm_personality mode requires a non-empty personal 3DMM history.") - with torch.no_grad(): - _, history_embedding = self.person_encoder(p) - person_embedding = self.personality_fusion(history_embedding, personality_embedding) - - if person_embedding.shape[0] != 1: - raise ValueError( - "ModifierNetwork expects exactly one personal reference per forward. " - "Use trainer-side micro-forwarding for batch training." - ) - return person_embedding - - def apply_weights(self, identity=False): - for idx, (name, module) in enumerate(self.hooked_modules.items()): - delta_w = torch.zeros_like(self.original_weights[idx]) if identity else self.kernel[idx] - if self.weight_kinds[idx] == "weight": - if self.hypernet_predict == "shift": - module.weight = self.original_weights[idx] + delta_w - elif self.hypernet_predict == "offset": - module.weight = self.original_weights[idx] * (self.tensor_1 + delta_w) - elif self.hypernet_predict == "weight": - module.weight = delta_w - else: - raise ValueError(f"Unsupported hypernet prediction mode: {self.hypernet_predict}") - else: - if "multihead_attn" in name and self.crossattn_modify == "kv": - delta_w = torch.cat((self.tensor_0, delta_w), dim=0) - if self.hypernet_predict == "shift": - module.in_proj_weight = self.original_weights[idx] + delta_w - elif self.hypernet_predict == "offset": - module.in_proj_weight = self.original_weights[idx] * (self.tensor_1 + delta_w) - elif self.hypernet_predict == "weight": - module.in_proj_weight = delta_w - else: - raise ValueError(f"Unsupported hypernet prediction mode: {self.hypernet_predict}") - - def forward(self, x, p=None, personality=None): - person_embedding = self.encode_person_condition(p=p, personality=personality) - self.kernel = self.hypernet(person_embedding) - self.apply_weights(identity=getattr(self, "_identity_modifier", False)) - # (P4) feed the same person embedding into the coarse-plan FiLM. Set it on - # the denoiser so every internal denoising step sees it; clear afterwards - # so a non-personalized path can never inherit a stale plan. - if self.person_coarse is not None: - p_gamma, p_beta = self.person_coarse(person_embedding) - self._coarse_denoiser._person_coarse_film = (p_gamma, p_beta) - try: - output = self.main_net(**x) - finally: - if self.person_coarse is not None: - self._coarse_denoiser._person_coarse_film = None - if self.regularization: - return output, self.regular_w * compute_regular_loss(self.kernel) - return output, person_embedding.new_tensor(0.0) diff --git a/personalised/code/framework/perfrdiff_rewrite_weight/person_specific/PersonSpecificEncoder.py b/personalised/code/framework/perfrdiff_rewrite_weight/person_specific/PersonSpecificEncoder.py deleted file mode 100644 index 68460ec32f2d4ef6db56a3cd6f02bdb9ee473031..0000000000000000000000000000000000000000 --- a/personalised/code/framework/perfrdiff_rewrite_weight/person_specific/PersonSpecificEncoder.py +++ /dev/null @@ -1,83 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class PositionalEncoding(nn.Module): - def __init__(self, d_model, max_len, device): - super().__init__() - self.encoding = torch.zeros(max_len, d_model, device=device) - self.encoding.requires_grad = False - pos = torch.arange(0, max_len, device=device).float().unsqueeze(1) - idx = torch.arange(0, d_model, step=2, device=device).float() - self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (idx / d_model))) - self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (idx / d_model))) - - def forward(self, x): - return self.encoding[:x.shape[1], :].unsqueeze(0) - - -class Transformer(nn.Module): - def __init__( - self, - device, - in_features, - embed_dim, - num_heads, - num_layers, - mlp_dim, - seq_len, - proj_dim, - proj_head="mlp", - drop_prob=0.1, - max_len=5000, - pos_encoding="absolute", - embed_layer="linear", - ): - super().__init__() - self.seq_len = seq_len - self.embed_dim = embed_dim if embed_layer == "linear" else in_features - self.embed_layer = nn.Linear(in_features, embed_dim) if embed_layer == "linear" else nn.Identity() - - if pos_encoding == "learnable": - self.pos_embed = nn.Parameter(torch.zeros(1, 1 + seq_len, self.embed_dim)) - elif pos_encoding == "absolute": - self.pos_embed = PositionalEncoding(d_model=self.embed_dim, max_len=max_len, device=device) - else: - raise NotImplementedError(f"Unsupported positional encoding: {pos_encoding}") - self.pos_encoding = pos_encoding - - self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim)) - encoder_layer = nn.TransformerEncoderLayer( - d_model=self.embed_dim, - nhead=num_heads, - dim_feedforward=mlp_dim, - batch_first=True, - ) - self.transformer = nn.TransformerEncoder(encoder_layer, num_layers) - self.dropout = nn.Dropout(p=drop_prob) - - if proj_head == "linear": - self.proj_head = nn.Linear(embed_dim, proj_dim) - elif proj_head == "mlp": - self.proj_head = nn.Sequential( - nn.Linear(embed_dim, embed_dim), - nn.ReLU(inplace=True), - nn.Linear(embed_dim, proj_dim), - ) - else: - self.proj_head = nn.Identity() - - def forward(self, x): - batch_size = x.shape[0] - x = self.embed_layer(x) - x = torch.cat([self.cls_token.expand(batch_size, -1, -1), x], dim=1) - x = x + (self.pos_embed(x) if self.pos_encoding == "absolute" else self.pos_embed) - x = self.dropout(x) - x = self.transformer(x) - feat = x[:, 0, :] - proj = F.normalize(self.proj_head(feat), dim=-1) - return feat, proj - - def get_model_name(self): - return self.__class__.__name__ diff --git a/personalised/code/framework/perfrdiff_rewrite_weight/person_specific/__init__.py b/personalised/code/framework/perfrdiff_rewrite_weight/person_specific/__init__.py deleted file mode 100644 index aa0217d1d639ad48f1b2eeab9ae00120cd71cccf..0000000000000000000000000000000000000000 --- a/personalised/code/framework/perfrdiff_rewrite_weight/person_specific/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .PersonSpecificEncoder import Transformer diff --git a/personalised/code/framework/utils/baseline_reaction_metrics.py b/personalised/code/framework/utils/baseline_reaction_metrics.py deleted file mode 100644 index 5a9edaa39015d0837e0feaa86deb4e8b56aceddf..0000000000000000000000000000000000000000 --- a/personalised/code/framework/utils/baseline_reaction_metrics.py +++ /dev/null @@ -1,407 +0,0 @@ -import argparse -import random -import time -import warnings -from collections import OrderedDict -from pathlib import Path - -import numpy as np -import torch -from tqdm import tqdm - -from framework.metrics import compute_FRC, compute_FRD, compute_FRVar, compute_TLCC, compute_s_mse -from framework.metrics.metric import baseline_mime, baseline_random -from framework.modules.post_processor import Processor -from framework.utils.compute_metrics import compute_MAE - - -BASELINE_NAMES = ("GT_identical", "B_Random", "B_Mime", "B_MeanFr") - - -def _repo_root(): - return Path(__file__).resolve().parents[2] - - -def _resolve_path(path): - path = Path(path).expanduser() - if path.is_absolute(): - return path - return (Path.cwd() / path).resolve() - - -def _is_sample_file(path): - return ( - path.suffix.lower() == ".npy" - and not path.name.startswith(".") - and not path.name.startswith("._") - ) - - -def _load_emotion(path): - try: - array = np.load(path, allow_pickle=False) - except ValueError as exc: - raise ValueError( - f"Cannot load numeric facial-attribute file: {path}. " - "It may be an AppleDouble/archive metadata file." - ) from exc - - tensor = torch.from_numpy(np.asarray(array)).float() - if tensor.dim() > 2: - tensor = tensor.squeeze() - if tensor.dim() == 1 and tensor.numel() == 25: - tensor = tensor.unsqueeze(0) - if tensor.dim() != 2 or tensor.shape[-1] != 25: - raise ValueError(f"Expected facial attributes with shape [T, 25], got {tuple(tensor.shape)} at {path}") - return tensor - - -def _iter_role_files(facial_dir, role): - role_dir = facial_dir / role - if not role_dir.is_dir(): - return - for session_dir in sorted(path for path in role_dir.iterdir() if path.is_dir() and not path.name.startswith(".")): - for path in sorted(session_dir.iterdir()): - if _is_sample_file(path): - yield Path(role) / session_dir.name / path.stem - - -def _build_gt_index(facial_dir): - index = {} - for role in ("speaker", "listener"): - for rel_path in _iter_role_files(facial_dir, role): - session_key = Path(rel_path.parts[0]) / rel_path.parts[1] - index.setdefault(session_key, []).append(rel_path) - return index - - -def _rel_to_path(facial_dir, rel_path): - return facial_dir / rel_path.with_suffix(".npy") - - -def _build_samples(facial_dir, num_preds, bidirectional, rng): - gt_index = _build_gt_index(facial_dir) - roles = ("speaker", "listener") if bidirectional else ("speaker",) - samples = [] - skipped = 0 - - for role in roles: - target_role = "listener" if role == "speaker" else "speaker" - for input_rel in _iter_role_files(facial_dir, role): - session = input_rel.parts[1] - target_rel = Path(target_role) / session / input_rel.name - if not _rel_to_path(facial_dir, target_rel).is_file(): - warnings.warn(f"Skip missing paired target: {target_rel}") - skipped += 1 - continue - - target_session = Path(target_role) / session - candidates = list(gt_index.get(target_session, [])) - extra_candidates = [path for path in candidates if path != target_rel] - extra_count = max(num_preds - 1, 0) - if extra_count == 0: - gt_rels = [target_rel] - elif len(extra_candidates) >= extra_count: - gt_rels = [target_rel] + rng.sample(extra_candidates, extra_count) - elif extra_candidates: - gt_rels = [target_rel] + rng.choices(extra_candidates, k=extra_count) - else: - gt_rels = [target_rel] + [target_rel] * extra_count - - samples.append( - { - "input": input_rel, - "target": target_rel, - "gts": gt_rels, - } - ) - - return samples, skipped - - -def _load_samples(facial_dir, samples, max_samples=None): - speaker_inputs = [] - raw_listener_targets = [] - sample_ids = [] - skipped = 0 - - selected = samples[:max_samples] if max_samples is not None else samples - for sample in selected: - try: - speaker_input = _load_emotion(_rel_to_path(facial_dir, sample["input"])) - targets = [_load_emotion(_rel_to_path(facial_dir, rel_path)) for rel_path in sample["gts"]] - except (OSError, ValueError) as exc: - warnings.warn(f"Skip invalid sample {sample['input']} -> {sample['target']}: {exc}") - skipped += 1 - continue - - speaker_inputs.append(speaker_input) - raw_listener_targets.append(targets) - sample_ids.append( - { - "input": str(sample["input"]), - "target": str(sample["target"]), - "gts": [str(path) for path in sample["gts"]], - } - ) - - return speaker_inputs, raw_listener_targets, sample_ids, skipped - - -def _compute_train_mean_fr(facial_dir): - listener_dir = facial_dir / "listener" - if not listener_dir.is_dir(): - raise FileNotFoundError(f"Missing training listener facial-attributes directory: {listener_dir}") - - sum_vector = torch.zeros(25, dtype=torch.float64) - frame_count = 0 - file_count = 0 - for rel_path in _iter_role_files(facial_dir, "listener"): - path = _rel_to_path(facial_dir, rel_path) - emotion = _load_emotion(path).double() - sum_vector += emotion.sum(dim=0) - frame_count += emotion.shape[0] - file_count += 1 - - if frame_count == 0: - raise RuntimeError(f"No listener facial-attribute frames found under {listener_dir}") - - print(f"Training listener mean uses {file_count} files and {frame_count} frames.") - return (sum_vector / frame_count).float() - - -def _build_target_alignment_predictions(speaker_inputs, num_preds): - return [ - speaker_input.new_zeros((num_preds, speaker_input.shape[0], speaker_input.shape[1])) - for speaker_input in speaker_inputs - ] - - -def _make_gt_identical(listener_targets): - return [ - target.clone() - for target in tqdm(listener_targets, desc="Building GT_identical", leave=False) - ] - - -def _make_random(listener_targets): - return [ - baseline_random(target) - for target in tqdm(listener_targets, desc="Building B_Random", leave=False) - ] - - -def _make_mime(speaker_inputs, num_preds): - predictions = [] - for speaker_input in tqdm(speaker_inputs, desc="Building B_Mime", leave=False): - if num_preds == 10: - predictions.append(baseline_mime(speaker_input)) - else: - predictions.append(speaker_input.unsqueeze(0).expand(num_preds, -1, -1).clone()) - return predictions - - -def _make_meanfr(speaker_inputs, train_mean_fr, num_preds): - mean_fr = train_mean_fr.view(1, 1, -1) - return [ - mean_fr.expand(num_preds, speaker_input.shape[0], -1).clone() - for speaker_input in tqdm(speaker_inputs, desc="Building B_MeanFr", leave=False) - ] - - -def _predictions_equal_targets(listener_predictions, listener_targets, atol=1e-6): - if len(listener_predictions) != len(listener_targets): - return False - - for prediction, target in zip(listener_predictions, listener_targets): - if prediction.shape != target.shape: - return False - if not torch.allclose(prediction, target, atol=atol, rtol=0.0): - return False - return True - - -def _identity_frc(listener_predictions): - num_preds = [prediction.shape[0] for prediction in listener_predictions] - return float(np.mean(num_preds)) - - -def compute_reaction_metrics( - speaker_inputs, - listener_predictions, - listener_targets, - threads=16, - desc="Metrics", - force_identity_frc=False): - metrics = OrderedDict( - [ - ("FRC", (compute_FRC, (listener_predictions, listener_targets), {"p": threads})), - ("FRD", (compute_FRD, (listener_predictions, listener_targets), {"p": threads})), - ("TLCC", (compute_TLCC, (listener_predictions, speaker_inputs), {"p": threads})), - ("smse", (compute_s_mse, (listener_predictions,), {})), - ("FRVar", (compute_FRVar, (listener_predictions,), {})), - ("MAE", (compute_MAE, (listener_predictions, listener_targets), {"p": threads})), - ] - ) - - results = {} - for name, (func, args, kwargs) in tqdm(metrics.items(), desc=desc): - t0 = time.perf_counter() - if name == "FRC" and force_identity_frc: - if _predictions_equal_targets(listener_predictions, listener_targets): - # The project FRC uses CCC. For exactly identical constant channels - # (common in binary AU streams), CCC is undefined and the metric code - # currently converts it to 0. GT_identical is an oracle baseline, so - # after verifying pred == target we report the expected perfect count. - value = _identity_frc(listener_predictions) - else: - warnings.warn( - "force_identity_frc=True but predictions differ from targets; " - "falling back to the standard FRC implementation." - ) - value = func(*args, **kwargs) - else: - value = func(*args, **kwargs) - if hasattr(value, "item"): - value = value.item() - elapsed = time.perf_counter() - t0 - results[name] = float(value) - results[f"{name}_time"] = elapsed - print(f"{name:6s} = {float(value):.6f} (time: {elapsed:.4f}s)") - return results - - -def _parse_args(): - parser = argparse.ArgumentParser( - description="Run no-training REACT facial-reaction baselines." - ) - parser.add_argument("--data_dir", required=True, help="REACT2025 data root containing train/val/test splits.") - parser.add_argument("--split", default="test", help="Evaluation split name.") - parser.add_argument("--train_split", default="train", help="Split used to compute B_MeanFr.") - parser.add_argument("--num_preds", type=int, default=10, help="Number of prediction samples per test item.") - parser.add_argument("--threads", type=int, default=16, help="Metric multiprocessing worker count.") - parser.add_argument("--seed", type=int, default=1234, help="Random seed for GT sampling and B_Random.") - parser.add_argument("--bidirectional", action="store_true", help="Evaluate both speaker->listener and listener->speaker.") - parser.add_argument("--max_samples", type=int, default=None, help="Optional smoke-test sample limit.") - parser.add_argument("--output_path", default="baseline_reaction_results.pt", help="Output .pt path.") - parser.add_argument("--post_config_name", default="configs/shared/model/emotion_autoencoder.yaml") - parser.add_argument("--post_ckpt_dir", default=None, help="Post-processor checkpoint directory.") - parser.add_argument("--post_clip_length", type=int, default=1000) - parser.add_argument("--device", default=None, help="Post-processor device, e.g. cuda:0 or cpu.") - return parser.parse_args() - - -def main(): - args = _parse_args() - if args.num_preds <= 1: - raise ValueError("--num_preds must be greater than 1 because smse requires multiple predictions.") - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - - repo_root = _repo_root() - data_dir = _resolve_path(args.data_dir) - eval_facial_dir = data_dir / args.split / "facial-attributes" - train_facial_dir = data_dir / args.train_split / "facial-attributes" - if not eval_facial_dir.is_dir(): - raise FileNotFoundError(f"Missing evaluation facial-attributes directory: {eval_facial_dir}") - if not train_facial_dir.is_dir(): - raise FileNotFoundError(f"Missing training facial-attributes directory: {train_facial_dir}") - - post_ckpt_dir = _resolve_path(args.post_ckpt_dir) if args.post_ckpt_dir else repo_root / "pretrained_models" / "post_processor" - if not (post_ckpt_dir / "checkpoint.pth").is_file(): - raise FileNotFoundError( - "Missing post-processor checkpoint. " - f"Expected: {post_ckpt_dir / 'checkpoint.pth'}" - ) - - rng = random.Random(args.seed) - samples, pair_skipped = _build_samples( - eval_facial_dir, - num_preds=args.num_preds, - bidirectional=args.bidirectional, - rng=rng, - ) - if not samples: - raise RuntimeError(f"No valid paired samples found under {eval_facial_dir}") - - speaker_inputs, raw_listener_targets, sample_ids, load_skipped = _load_samples( - eval_facial_dir, - samples, - max_samples=args.max_samples, - ) - if not speaker_inputs: - raise RuntimeError("No samples could be loaded after filtering invalid files.") - - print( - f"Loaded {len(speaker_inputs)} samples from split={args.split}; " - f"pair_skipped={pair_skipped}; load_skipped={load_skipped}; " - f"bidirectional={args.bidirectional}." - ) - - train_mean_fr = _compute_train_mean_fr(train_facial_dir) - - device = torch.device(args.device) if args.device else torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - post_processor = Processor( - config_name=args.post_config_name, - ckpt_dir=str(post_ckpt_dir), - cfg_dir=str(repo_root), - clip_len_test=args.post_clip_length, - device=device, - num_preds=args.num_preds, - ) - target_alignment_predictions = _build_target_alignment_predictions(speaker_inputs, args.num_preds) - listener_targets = post_processor.forward( - prediction_list=target_alignment_predictions, - target_list=raw_listener_targets, - ) - - baseline_predictions = OrderedDict( - [ - ("GT_identical", _make_gt_identical(listener_targets)), - ("B_Random", _make_random(listener_targets)), - ("B_Mime", _make_mime(speaker_inputs, args.num_preds)), - ("B_MeanFr", _make_meanfr(speaker_inputs, train_mean_fr, args.num_preds)), - ] - ) - - metrics = OrderedDict() - for name, predictions in baseline_predictions.items(): - print(f"\n=== {name} ===") - metrics[name] = compute_reaction_metrics( - speaker_inputs=speaker_inputs, - listener_predictions=predictions, - listener_targets=listener_targets, - threads=args.threads, - desc=f"Evaluating {name}", - force_identity_frc=(name == "GT_identical"), - ) - - output = { - "GT": listener_targets, - "INPUT": speaker_inputs, - "PRED": baseline_predictions, - "metrics": metrics, - "train_mean_fr": train_mean_fr, - "sample_ids": sample_ids, - "config": { - "data_dir": str(data_dir), - "split": args.split, - "train_split": args.train_split, - "num_preds": args.num_preds, - "bidirectional": args.bidirectional, - "seed": args.seed, - "threads": args.threads, - "max_samples": args.max_samples, - }, - } - - output_path = _resolve_path(args.output_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - torch.save(output, output_path) - print(f"\nSaved baseline results to {output_path}") - - -if __name__ == "__main__": - main() diff --git a/personalised/code/framework/utils/compute_metrics.py b/personalised/code/framework/utils/compute_metrics.py deleted file mode 100644 index 9ccbd176117c3b7c58fafcfdd544f45db02894d5..0000000000000000000000000000000000000000 --- a/personalised/code/framework/utils/compute_metrics.py +++ /dev/null @@ -1,271 +0,0 @@ -import time -import numpy as np -import multiprocessing as mp -import torch -from framework.metrics import * -from tslearn.metrics import dtw - - -def _func(target, pred): - # target: (10, l, dim) - # pred: (10, l, dim) - num_preds = pred.shape[0] - mean_mae_sum = 0 - for i in range(num_preds): - mae_list = [] - for j in range(num_preds): - mae = np.mean(np.abs(target[j][:, 15:].numpy() - pred[i][:, 15:].numpy())) - mae_list.append(mae) - mean_mae_sum += np.mean(mae_list) - return mean_mae_sum / num_preds - - -def compute_MAE(preds, targets, p=4): - MAE_list = [] - with mp.Pool(processes=p) as pool: - MAE_list += pool.starmap(_func, zip(targets, preds)) - return np.mean(MAE_list) - - -def compute_metrics(speaker_inputs, - listener_predictions, - listener_targets, - threads=16): - - metrics = { - 'FRC': (compute_FRC, (listener_predictions, listener_targets), {'p': threads}), - 'FRD': (compute_FRD, (listener_predictions, listener_targets), {'p': threads}), - 'TLCC': (compute_TLCC, (listener_predictions, speaker_inputs), {'p': threads}), - 'smse': (compute_s_mse, (listener_predictions,), {}), - 'FRVar': (compute_FRVar, (listener_predictions,), {}), - # 'MAE': (compute_MAE, (listener_predictions, listener_targets), {'p': threads}), - } - - results = {} - for name, (func, args, kwargs) in metrics.items(): - t0 = time.perf_counter() - value = func(*args, **kwargs) - if hasattr(value, 'item'): - value = value.item() - elapsed = time.perf_counter() - t0 - - results[name] = value - results[f"{name}_time"] = elapsed - print(f"{name:6s} = {value:.6f} (time: {elapsed:.4f}s)") - - return results - - -def _to_3d_tensor(item): - if isinstance(item, list): - item = torch.stack(item, dim=0) - item = torch.as_tensor(item).float() - if item.dim() == 2: - item = item.unsqueeze(0) - return item - - -def _eeg_mask_for(mask, target_idx, length, dim): - if mask is None: - return torch.ones(length, dim) - mask = _to_3d_tensor(mask) - if mask.shape[0] == 1: - mask_item = mask[0] - else: - mask_item = mask[min(target_idx, mask.shape[0] - 1)] - return mask_item[:length, :dim].float() - - -def _masked_ccc_1d(target, prediction, mask): - valid = mask > 0.5 - valid = valid & torch.isfinite(target) & torch.isfinite(prediction) - if valid.sum().item() < 2: - return None - y_true = target[valid].detach().cpu().numpy() - y_pred = prediction[valid].detach().cpu().numpy() - std_true = np.std(y_true) - std_pred = np.std(y_pred) - if std_true < 1e-8 or std_pred < 1e-8: - return 0.0 - cor = np.corrcoef(y_true, y_pred)[0, 1] - mean_true = np.mean(y_true) - mean_pred = np.mean(y_pred) - var_true = np.var(y_true) - var_pred = np.var(y_pred) - return (2 * cor * std_true * std_pred) / (var_true + var_pred + (mean_true - mean_pred) ** 2 + 1e-8) - - -def _masked_channel_mean_ccc(target, prediction, mask): - scores = [] - for dim_idx in range(target.shape[-1]): - score = _masked_ccc_1d(target[:, dim_idx], prediction[:, dim_idx], mask[:, dim_idx]) - if score is not None and np.isfinite(score): - scores.append(score) - return float(np.mean(scores)) if scores else None - - -def _masked_channel_mean_dtw(target, prediction, mask): - scores = [] - for dim_idx in range(target.shape[-1]): - valid = mask[:, dim_idx] > 0.5 - valid = valid & torch.isfinite(target[:, dim_idx]) & torch.isfinite(prediction[:, dim_idx]) - if valid.sum().item() < 1: - continue - target_ch = target[valid, dim_idx].detach().cpu().numpy().astype(np.float32) - pred_ch = prediction[valid, dim_idx].detach().cpu().numpy().astype(np.float32) - scores.append(dtw(pred_ch, target_ch)) - return float(np.mean(scores)) if scores else None - - -def _compute_eeg_frc(preds, targets, masks): - sample_scores = [] - for sample_idx, (prediction, target) in enumerate(zip(preds, targets)): - prediction = _to_3d_tensor(prediction) - target = _to_3d_tensor(target) - mask = masks[sample_idx] if masks is not None else None - length = min(prediction.shape[1], target.shape[1]) - dim = min(prediction.shape[2], target.shape[2]) - prediction = prediction[:, :length, :dim] - target = target[:, :length, :dim] - - pred_scores = [] - for pred_idx in range(prediction.shape[0]): - target_scores = [] - for target_idx in range(target.shape[0]): - target_mask = _eeg_mask_for(mask, target_idx, length, dim) - score = _masked_channel_mean_ccc(target[target_idx], prediction[pred_idx], target_mask) - if score is not None: - target_scores.append(score) - if target_scores: - pred_scores.append(max(target_scores)) - if pred_scores: - sample_scores.append(float(np.sum(pred_scores))) - return float(np.mean(sample_scores)) if sample_scores else float("nan") - - -def _compute_eeg_frd(preds, targets, masks): - sample_scores = [] - for sample_idx, (prediction, target) in enumerate(zip(preds, targets)): - prediction = _to_3d_tensor(prediction) - target = _to_3d_tensor(target) - mask = masks[sample_idx] if masks is not None else None - length = min(prediction.shape[1], target.shape[1]) - dim = min(prediction.shape[2], target.shape[2]) - prediction = prediction[:, :length, :dim] - target = target[:, :length, :dim] - - pred_scores = [] - for pred_idx in range(prediction.shape[0]): - target_scores = [] - for target_idx in range(target.shape[0]): - target_mask = _eeg_mask_for(mask, target_idx, length, dim) - score = _masked_channel_mean_dtw(target[target_idx], prediction[pred_idx], target_mask) - if score is not None: - target_scores.append(score) - if target_scores: - pred_scores.append(min(target_scores)) - if pred_scores: - sample_scores.append(float(np.sum(pred_scores))) - return float(np.mean(sample_scores)) if sample_scores else float("nan") - - -def _compute_eeg_tlcc(preds, targets, masks, seconds=2, fps=1): - sample_offsets = [] - max_lag = int(seconds * fps) - for sample_idx, (prediction, target) in enumerate(zip(preds, targets)): - prediction = _to_3d_tensor(prediction) - target = _to_3d_tensor(target) - mask = masks[sample_idx] if masks is not None else None - length = min(prediction.shape[1], target.shape[1]) - dim = min(prediction.shape[2], target.shape[2]) - prediction = prediction[:, :length, :dim] - target = target[0, :length, :dim] - target_mask = _eeg_mask_for(mask, 0, length, dim) - - pred_offsets = [] - for pred_idx in range(prediction.shape[0]): - channel_offsets = [] - for dim_idx in range(dim): - lag_scores = [] - for lag in range(-max_lag, max_lag + 1): - if lag > 0: - pred_ch = prediction[pred_idx, lag:, dim_idx] - target_ch = target[:-lag, dim_idx] - mask_ch = target_mask[:-lag, dim_idx] - elif lag < 0: - pred_ch = prediction[pred_idx, :lag, dim_idx] - target_ch = target[-lag:, dim_idx] - mask_ch = target_mask[-lag:, dim_idx] - else: - pred_ch = prediction[pred_idx, :, dim_idx] - target_ch = target[:, dim_idx] - mask_ch = target_mask[:, dim_idx] - ccc = _masked_ccc_1d(target_ch, pred_ch, mask_ch) - if ccc is not None: - lag_scores.append((ccc, lag)) - if lag_scores: - best_lag = max(lag_scores, key=lambda item: item[0])[1] - channel_offsets.append(abs(best_lag)) - if channel_offsets: - pred_offsets.append(float(np.mean(channel_offsets))) - if pred_offsets: - sample_offsets.append(float(np.mean(pred_offsets))) - return float(np.mean(sample_offsets)) if sample_offsets else float("nan") - - -def _compute_eeg_smse(preds): - distances = [] - for prediction in preds: - prediction = _to_3d_tensor(prediction) - if prediction.shape[0] < 2: - distances.append(prediction.new_tensor(0.0)) - continue - flattened = prediction.reshape(prediction.shape[0], -1) - dist = torch.pow(torch.cdist(flattened, flattened), 2) - dist = torch.sum(dist) / (flattened.shape[0] * (flattened.shape[0] - 1) * flattened.shape[1]) - distances.append(dist) - return torch.stack(distances).mean().item() if distances else float("nan") - - -def _compute_eeg_frvar(preds, masks=None): - values = [] - for sample_idx, prediction in enumerate(preds): - prediction = _to_3d_tensor(prediction) - if masks is None: - values.append(torch.var(prediction, dim=1, unbiased=False).mean()) - continue - mask = _to_3d_tensor(masks[sample_idx])[0] - length = min(prediction.shape[1], mask.shape[0]) - dim = min(prediction.shape[2], mask.shape[1]) - valid = mask[:length, :dim] > 0.5 - if valid.sum().item() == 0: - continue - var = torch.var(prediction[:, :length, :dim], dim=1, unbiased=False) - channel_valid = valid.any(dim=0) - if channel_valid.any(): - values.append(var[:, channel_valid].mean()) - return torch.stack(values).mean().item() if values else float("nan") - - -def compute_eeg_metrics(listener_predictions, - listener_targets, - masks=None, - threads=16): - metrics = { - "EEG_FRC": lambda: _compute_eeg_frc(listener_predictions, listener_targets, masks), - "EEG_FRD": lambda: _compute_eeg_frd(listener_predictions, listener_targets, masks), - "EEG_TLCC": lambda: _compute_eeg_tlcc(listener_predictions, listener_targets, masks), - "EEG_smse": lambda: _compute_eeg_smse(listener_predictions), - "EEG_FRVar": lambda: _compute_eeg_frvar(listener_predictions, masks), - } - results = {} - for name, func in metrics.items(): - t0 = time.perf_counter() - value = func() - elapsed = time.perf_counter() - t0 - if hasattr(value, "item"): - value = value.item() - results[name] = value - results[f"{name}_time"] = elapsed - print(f"{name:10s} = {value:.6f} (time: {elapsed:.4f}s)") - return results diff --git a/personalised/code/framework/utils/losses.py b/personalised/code/framework/utils/losses.py deleted file mode 100644 index 521f26361a4c73687c90463dfd9962dfdafa88d6..0000000000000000000000000000000000000000 --- a/personalised/code/framework/utils/losses.py +++ /dev/null @@ -1,298 +0,0 @@ -from __future__ import print_function - -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange - - -class KLLoss(nn.Module): - def __init__(self): - super(KLLoss, self).__init__() - - def forward(self, q, p): - div = torch.distributions.kl_divergence(q, p) - return div.mean() - - def __repr__(self): - return "KLLoss()" - - -class VAELoss(nn.Module): - def __init__(self, kl_p: float = 0.0002, - w_emo: float = None, - w_exp: float = None, - w_rot: float = None, - w_tran: float = None, - eeg_loss_weight: float = 1.0, - **kwargs): - super(VAELoss, self).__init__() - self.mse = nn.MSELoss(reduce=True, size_average=True) - self.kl_loss = KLLoss() - self.kl_p = kl_p - self.eeg_loss_weight = eeg_loss_weight - - if w_emo is None: - w_emo = 1.0 - if w_exp is None: - w_exp = 1.0 - if w_rot is None: - w_rot = 10.0 - if w_tran is None: - w_tran = 10.0 - self.w_emo = w_emo - self.w_exp = w_exp - self.w_rot = w_rot - self.w_tran = w_tran - - @staticmethod - def masked_mse(prediction, target, mask): - mask = mask.to(dtype=prediction.dtype) - loss = ((prediction - target) ** 2) * mask - return loss.sum() / mask.sum().clamp_min(1.0) - - def forward(self, gt_emotions, gt_3dmms, pred_emotions, pred_3dmms, distribution, - prediction_eeg=None, target_eeg=None, target_eeg_mask=None): - """ List - gt_emotion; gt_3dmm; pred_emotion; pred_3dmm - """ - bsz = len(gt_emotions) - rec_emotion_loss = 0 - rec_param_loss = 0 - for gt_emotion, gt_3dmm, pred_emotion, pred_3dmm in zip( - gt_emotions, gt_3dmms, pred_emotions, pred_3dmms): - gt_emotion = gt_emotion.to(pred_emotion.get_device()) - gt_3dmm = gt_3dmm.to(pred_3dmm.get_device()) - - exp_part = self.w_exp * self.mse(pred_3dmm[:, :52], gt_3dmm[:, :52]) # expression - rot_part = self.w_rot * self.mse(pred_3dmm[:, 52:55], gt_3dmm[:, 52:55]) # rotation - tran_part = self.w_tran * self.mse(pred_3dmm[:, 55:], gt_3dmm[:, 55:]) # translation - rec_param_loss = rec_param_loss + (exp_part + rot_part + tran_part) - rec_emotion_loss = rec_emotion_loss + self.w_emo * self.mse(pred_emotion, gt_emotion) - - rec_emotion_loss = rec_emotion_loss / bsz - rec_param_loss = rec_param_loss / bsz - rec_loss = rec_emotion_loss + rec_param_loss - - mu_ref = torch.zeros_like(distribution[0].loc).to(gt_emotion.get_device()) - scale_ref = torch.ones_like(distribution[0].scale).to(gt_emotion.get_device()) - distribution_ref = torch.distributions.Normal(mu_ref, scale_ref) - - kld_loss = 0 - for t in range(len(distribution)): - kld_loss += self.kl_loss(distribution[t], distribution_ref) - kld_loss = kld_loss / len(distribution) - - loss = rec_loss + self.kl_p * kld_loss - loss_eeg = loss.new_tensor(0.0) - eeg_valid_ratio = loss.new_tensor(0.0) - if prediction_eeg is not None and target_eeg is not None: - target_eeg = target_eeg.to(prediction_eeg.device).float() - target_eeg_mask = target_eeg_mask.to(prediction_eeg.device).float() \ - if target_eeg_mask is not None else torch.ones_like(target_eeg) - loss_eeg = self.masked_mse(prediction_eeg, target_eeg, target_eeg_mask) - eeg_valid_ratio = target_eeg_mask.float().mean() - loss = loss + self.eeg_loss_weight * loss_eeg - return loss, rec_loss, rec_emotion_loss, rec_param_loss, kld_loss, loss_eeg, eeg_valid_ratio - - def __repr__(self): - return "VAELoss()" - - -def div_loss(Y_1_list, Y_2_list): - loss = 0.0 - B = len(Y_1_list) - for y1, y2 in zip(Y_1_list, Y_2_list): - y1_flat = y1.view(1, -1) - y2_flat = y2.view(1, -1) - Y = torch.cat([y1_flat, y2_flat], dim=0) - dist2 = F.pdist(Y, 2) ** 2 - loss += (-dist2 / 100).exp().mean() - loss = loss / B - return loss - - -def div_loss_v2(Y_1, Y_2): - loss = 0.0 - b,t,c = Y_1.shape - Y_g = torch.cat([Y_1.view(b,1,-1), Y_2.view(b,1,-1)], dim = 1) - for Y in Y_g: - dist = F.pdist(Y, 2) ** 2 - loss += (-dist / 100).exp().mean() - loss /= b - return loss - - -def TemporalLoss(Y): - diff = Y[:, 1:, :] - Y[:, :-1, :] - t_loss = torch.mean(torch.norm(diff, dim=2, p=2) ** 2) - return t_loss - - -def L1Loss(prediction, target, reduction="min", **kwargs): - # prediction has shape of [batch_size, num_preds, features] - # target has shape of [batch_size, num_preds, features] - assert len(prediction.shape) == len(target.shape), "prediction and target must have the same shape" - assert len(prediction.shape) == 3, "Only works with predictions of shape [batch_size, num_preds, features]" - - # manual implementation of L1 loss - loss = (torch.abs(prediction - target)).mean(dim=-1) - - # reduce across multiple predictions - if reduction == "mean": - loss = torch.mean(loss) - elif reduction == "min": - loss = loss.min(dim=-1)[0].mean() - else: - raise NotImplementedError("reduction {} not implemented".format(reduction)) - return loss - - -def MSELoss(prediction, target, reduction="mean", **kwargs): - assert len(prediction.shape) == len(target.shape), "prediction and target must have the same shape" - assert len(prediction.shape) == 3, "Only works with predictions of shape [batch_size, num_preds, features]" - - loss = ((prediction - target) ** 2).mean(dim=-1) - - # reduce across multiple predictions - if reduction == "mean": - loss = torch.mean(loss) - elif reduction == "min": - loss = loss.min(dim=-1)[0].mean() - else: - raise NotImplementedError("reduction {} not implemented".format(reduction)) - return loss - - -def MSELossApt(prediction, target, reduction="mean", - w_au=1.0, w_va=5.0, w_em=2.0, **kwargs): - assert len(prediction.shape) == 3, "Only works with predictions of shape [batch_size, num_preds, features]" - loss_au = F.mse_loss(prediction[:, :, :15], target[:, :, :15], reduction=reduction) - loss_va = F.mse_loss(prediction[:, :, 15:17], target[:, :, 15:17], reduction=reduction) - loss_em = F.mse_loss(prediction[:, :, 17:], target[:, :, 17:], reduction=reduction) - loss = loss_au * w_au + loss_va * w_va + loss_em * w_em - - losses_dict = {"loss": loss, "loss_au": loss_au, "loss_va": loss_va, "loss_em": loss_em} - return losses_dict - - -class DiffusionLoss: - def __init__(self, - losses_type='MSELoss', - n_preds=10, - prior_loss_weight=1.0, - eeg_loss_weight=1.0, - w_au=1.0, # action unit - w_va=5.0, # valence and arousal - w_em=2.0, # emotion - **kwargs): - self.loss_type = losses_type - self.n_preds = n_preds - self.prior_loss_weight = prior_loss_weight - self.eeg_loss_weight = eeg_loss_weight - self.w_au = w_au - self.w_va = w_va - self.w_em = w_em - - @staticmethod - def masked_mse(prediction, target, mask): - mask = mask.to(dtype=prediction.dtype) - loss = ((prediction - target) ** 2) * mask - return loss.sum() / mask.sum().clamp_min(1.0) - - def __call__(self, output_decoder): - output_prior = None - if "output_decoder" in output_decoder: - output_prior = output_decoder.get("output_prior") - output_decoder = output_decoder["output_decoder"] - - prediction_emotion = output_decoder["prediction_emotion"] - target_emotion = output_decoder["target_emotion"] - - _, _, window_size, emotion_dim = prediction_emotion.shape - prediction_emotion = prediction_emotion.reshape(-1, self.n_preds, window_size * emotion_dim) - target_emotion = target_emotion.reshape(-1, self.n_preds, window_size * emotion_dim) - losses_dict = eval(self.loss_type)( - prediction_emotion, target_emotion, k=self.n_preds, w_au=self.w_au, w_va=self.w_va, w_em=self.w_em) - - losses_dict["loss_eeg"] = losses_dict["loss"].new_tensor(0.0) - losses_dict["eeg_valid_ratio"] = losses_dict["loss"].new_tensor(0.0) - if "prediction_eeg" in output_decoder and "target_eeg" in output_decoder: - prediction_eeg = output_decoder["prediction_eeg"] - target_eeg = output_decoder["target_eeg"] - target_eeg_mask = output_decoder.get("target_eeg_mask", torch.ones_like(target_eeg)) - loss_eeg = self.masked_mse(prediction_eeg, target_eeg, target_eeg_mask) - losses_dict["loss_eeg"] = loss_eeg - losses_dict["eeg_valid_ratio"] = target_eeg_mask.float().mean() - losses_dict["loss"] = losses_dict["loss"] + self.eeg_loss_weight * loss_eeg - - if output_prior is not None: - encoded_prediction = output_prior["encoded_prediction"] - encoded_target = output_prior["encoded_target"] - if encoded_prediction.dim() == 4: - encoded_prediction = encoded_prediction.squeeze(-2) - if encoded_target.dim() == 4: - encoded_target = encoded_target.squeeze(-2) - loss_prior = MSELoss(encoded_prediction, encoded_target, reduction="mean") - losses_dict["loss_prior"] = loss_prior - losses_dict["loss"] = losses_dict["loss"] + self.prior_loss_weight * loss_prior - else: - losses_dict["loss_prior"] = losses_dict["loss"].new_tensor(0.0) - - return losses_dict - -class EmotionVAELoss: - def __init__(self, w_au, w_va, w_em, w_kld, **kwargs): - self.w_au = w_au - self.w_va = w_va - self.w_em = w_em - self.w_kld = w_kld - - self.au_criterion = nn.BCEWithLogitsLoss(reduction="none") # "mean" - self.va_criterion = nn.MSELoss(reduction="mean") - self.em_criterion = nn.KLDivLoss(reduction="none") # batchmean - self.kld_criterion = KLLoss() - - def kld_loss(self, mu, logvar, reduction='batchmean'): - kld_element = 1 + logvar - mu.pow(2) - logvar.exp() - if reduction == 'sum': - # sum over B, L, D - return -0.5 * torch.sum(kld_element) - elif reduction == 'batchmean': - # sum over L,D, then mean over batch - kld = -0.5 * torch.sum(kld_element, dim=(1, 2)) # [B] - return torch.mean(kld) # scalar - elif reduction == 'mean': - # mean over all elements - return -0.5 * torch.mean(kld_element) - else: - raise ValueError("Unknown reduction") - - def __call__(self, predictions, targets, distribution, mask=None): - au_logits, va_logits, emotion_logits = predictions - au_targets, va_targets, emotion_targets = \ - targets[:, :, :15], targets[:, :, 15:17], targets[:, :, 17:] - - au_loss = self.au_criterion(au_logits, au_targets) # AUs - va_loss = self.va_criterion(va_logits, va_targets) # valence and arousal - emotion_logits = rearrange(F.log_softmax(emotion_logits, dim=-1), "b l d -> (b l) d") - emotion_targets = rearrange(emotion_targets, "b l d -> (b l) d") - em_loss = self.em_criterion(emotion_logits, emotion_targets) - - if mask is not None: - au_mask = mask # [B, L, 1] - au_loss = (au_loss * au_mask).sum() / (au_mask.sum() * au_logits.shape[-1]) - em_mask = rearrange(mask, "b l d -> (b l) d") # [BxL, 1] - em_loss = (em_loss * em_mask).sum() / em_mask.sum() - - mu_ref = torch.zeros_like(distribution.loc).to(au_logits) - scale_ref = torch.ones_like(distribution.scale).to(au_logits) - distribution_ref = torch.distributions.Normal(mu_ref, scale_ref) - kld_loss = self.kld_criterion(distribution, distribution_ref) - - loss = self.w_kld * kld_loss + \ - self.w_au * au_loss + \ - self.w_va * va_loss + \ - self.w_em * em_loss - - return loss, kld_loss, au_loss, va_loss, em_loss diff --git a/personalised/code/framework/utils/losses_causal.py b/personalised/code/framework/utils/losses_causal.py deleted file mode 100644 index 67a71c656ac2596d81cec4dea16da395ec9ac8e1..0000000000000000000000000000000000000000 --- a/personalised/code/framework/utils/losses_causal.py +++ /dev/null @@ -1,52 +0,0 @@ -""" -DiffusionLossCoarse -=================== -Adds an explicit coarse-to-fine cross-entropy term on top of the baseline -`DiffusionLoss`. The coarse head predicts the listener's 8-class facial -expression distribution per timestep; we supervise it with a masked soft -cross-entropy against the ground-truth distribution (target_emotion[..., 17:25]). - -Padded / invalid frames have an all-zero target distribution and are masked out -(soft-CE against an all-zero label is naturally zero, but we additionally -renormalise valid frames to a proper distribution). -""" -import torch -import torch.nn.functional as F - -from framework.utils.losses import DiffusionLoss - - -class DiffusionLossCoarse(DiffusionLoss): - def __init__(self, w_coarse: float = 0.5, coarse_emo_start: int = 17, - coarse_classes: int = 8, **kwargs): - super().__init__(**kwargs) - self.w_coarse = float(w_coarse) - self.coarse_emo_start = int(coarse_emo_start) - self.coarse_classes = int(coarse_classes) - - def __call__(self, output_decoder): - losses = super().__call__(output_decoder) - - dec = output_decoder.get("output_decoder", output_decoder) \ - if isinstance(output_decoder, dict) else output_decoder - - coarse_logits = dec.get("coarse_logits") if isinstance(dec, dict) else None - if coarse_logits is None: - losses["loss_coarse"] = losses["loss"].new_tensor(0.0) - return losses - - target_emotion = dec["target_emotion"] # (bs, num_preds, T, 25) - s0 = self.coarse_emo_start - s1 = s0 + self.coarse_classes - tgt = target_emotion[..., s0:s1] # (bs, np, T, C) - - denom = tgt.sum(dim=-1, keepdim=True) # (bs, np, T, 1) - valid = (denom > 0.5).to(coarse_logits.dtype) # real-distribution frames - prob = tgt / denom.clamp_min(1e-6) # normalised soft labels - logp = F.log_softmax(coarse_logits, dim=-1) - ce = -(prob * logp).sum(dim=-1, keepdim=True) # (bs, np, T, 1) - loss_coarse = (ce * valid).sum() / valid.sum().clamp_min(1.0) - - losses["loss_coarse"] = loss_coarse - losses["loss"] = losses["loss"] + self.w_coarse * loss_coarse - return losses diff --git a/personalised/code/framework/utils/util.py b/personalised/code/framework/utils/util.py deleted file mode 100644 index 4443d42406a9a2c627c34f7299a355966c47be97..0000000000000000000000000000000000000000 --- a/personalised/code/framework/utils/util.py +++ /dev/null @@ -1,139 +0,0 @@ -import os -import numpy as np -import torch -from datetime import datetime -import yaml -from sklearn.manifold import TSNE -import matplotlib.pyplot as plt -from omegaconf import OmegaConf -from torch.backends import cudnn -import hydra - - -def init_seed(seed, rank=0): - process_seed = seed + rank - torch.manual_seed(process_seed) - torch.cuda.manual_seed(process_seed) - np.random.seed(process_seed) - cudnn.benchmark = True - torch.backends.cudnn.benchmark = True - torch.backends.cudnn.deterministic = False - - -def load_config_from_file(path): - return OmegaConf.load(path) - - -def load_config(args=None, config_path=None): - if args is not None: - config_from_args = OmegaConf.create(vars(args)) - else: - config_from_args = OmegaConf.from_cli() - # config_from_file = OmegaConf.load(cli_conf.pop('config') if config_path is None else config_path) - config_from_file = load_config_from_file(config_path) - return OmegaConf.merge(config_from_file, config_from_args) - - -def store_config(config): - dir = config.trainer.out_dir - os.makedirs(dir, exist_ok=True) - with open(os.path.join(dir, "config.yaml"), "w") as f: - yaml.dump(OmegaConf.to_container(config), f) - - -def torch_img_to_np(img): - return img.detach().cpu().numpy().transpose(0, 2, 3, 1) - - -def torch_img_to_np2(img): - img = img.detach().cpu().numpy() - img = img * np.array([0.5, 0.5, 0.5]).reshape(1, -1, 1, 1) - img = img + np.array([0.5, 0.5, 0.5]).reshape(1, -1, 1, 1) - img = img.transpose(0, 2, 3, 1) - img = img * 255.0 - img = np.clip(img, 0, 255).astype(np.uint8)[:, :, :, [2, 1, 0]] - - return img - - -def _fix_image(image): - if image.max() < 30.: - image = image * 255. - image = np.clip(image, 0, 255).astype(np.uint8)[:, :, :, [2, 1, 0]] - return image - - -class AverageMeter(object): - def __init__(self): - self.reset() - - def reset(self): - self.val = 0 - self.avg = 0 - self.sum = 0 - self.count = 0 - - def update(self, val, n=1): - self.val = val - self.sum += val * n - self.count += n - self.avg = self.sum / self.count - - -def collect_grad_value_(parameters): - grad_values = [] - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - for p in filter(lambda p: p.grad is not None, parameters): - grad_values.append(p.grad.data.abs().mean().item()) - grad_values = np.array(grad_values) - return grad_values - - -def save_checkpoint(checkpoint_path, net, optimizer=None, epoch=None, best_loss=None): - checkpoint = { - 'epoch': epoch if epoch is not None else None, - 'best_loss': best_loss if best_loss is not None else None, - 'state_dict': net.state_dict(), - 'optimizer': optimizer.state_dict() if optimizer is not None else None, - } - torch.save(checkpoint, checkpoint_path) - - -def from_pretrained_checkpoint(checkpoint_path, model, device): - checkpoint = torch.load(checkpoint_path, map_location='cpu') - if isinstance(model, torch.optim.Optimizer): - model.load_state_dict(checkpoint['optimizer']) - print(f'Successfully load optimizer checkpoint: {checkpoint_path}') - else: - model.load_state_dict(checkpoint['state_dict']) - model.to(device) - print(f'Successfully load model checkpoint: {checkpoint_path}') - return checkpoint.get('best_loss', float('inf')), checkpoint.get('epoch', 0) - - -def get_lr(optimizer): - for param_group in optimizer.param_groups: - return param_group['lr'] - - -def collect_grad_stats(parameters): - grad_values = [] - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - for p in filter(lambda p: p.grad is not None, parameters): - # Store the absolute values of gradients - grad_values.extend(p.grad.data.abs().view(-1).cpu().numpy()) - - # Convert to a numpy array for statistical computation - grad_values = np.array(grad_values) - if grad_values.size == 0: - return {"min": None, "max": None, "mean": None} - - # Compute min, max, and mean - grad_stats = { - "min": grad_values.min(), - "max": grad_values.max(), - "mean": grad_values.mean() - } - return grad_stats diff --git a/personalised/code/launch/__init__.py b/personalised/code/launch/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/personalised/code/launch/blender.py b/personalised/code/launch/blender.py deleted file mode 100644 index 4ccf32cd715bbd2c8e355f9ccbc39b13f8481872..0000000000000000000000000000000000000000 --- a/personalised/code/launch/blender.py +++ /dev/null @@ -1,27 +0,0 @@ -# Fix blender path -import sys -import os - -# local packages -# sys.path.append(os.path.expanduser("~/.local/lib/python3.9/site-packages")) -sys.path.insert(0, os.path.expanduser("~/.local/lib/python3.9/site-packages")) -import bpy -import os -from argparse import ArgumentParser - - -# Monkey patch argparse such that -# blender / python / hydra parsing works -def parse_args(self, args=None, namespace=None): - if args is not None: - return self.parse_args_bak(args=args, namespace=namespace) - try: - idx = sys.argv.index("--") - args = sys.argv[idx + 1 :] # the list after '--' - except ValueError as e: # '--' not in the list: - args = [] - return self.parse_args_bak(args=args, namespace=namespace) - - -setattr(ArgumentParser, "parse_args_bak", ArgumentParser.parse_args) -setattr(ArgumentParser, "parse_args", parse_args) diff --git a/personalised/code/launch/prepare.py b/personalised/code/launch/prepare.py deleted file mode 100644 index 617558ed2a849598d4afa20b08442d94fca02794..0000000000000000000000000000000000000000 --- a/personalised/code/launch/prepare.py +++ /dev/null @@ -1,75 +0,0 @@ -import os -import warnings -from pathlib import Path -from omegaconf import OmegaConf -from utils.runid import generate_id -import hydra -import torch - - -# Local paths -def code_path(path=""): - code_dir = hydra.utils.get_original_cwd() - code_dir = Path(code_dir) - return str(code_dir / path) - - -def working_path(path): - return str(Path(os.getcwd()) / path) - - -# fix the id for this run -ID = generate_id() - - -def generate_id(): - return ID - - -def get_last_checkpoint(path, version=None): - # ckpt_name="last.ckpt" - output_dir = Path(hydra.utils.to_absolute_path(path)) - if version is None: - last_ckpt_path = output_dir / "checkpoints" - else: - last_ckpt_path = output_dir / f"checkpoints/{version}" - return str(last_ckpt_path) - - -OmegaConf.register_new_resolver("code_path", code_path) -OmegaConf.register_new_resolver("working_path", working_path) -OmegaConf.register_new_resolver("generate_id", generate_id) -OmegaConf.register_new_resolver("absolute_path", hydra.utils.to_absolute_path) -OmegaConf.register_new_resolver("get_last_checkpoint", get_last_checkpoint) - - -# Remove some warnings -warnings.filterwarnings( - "ignore", ".*Trying to infer the `batch_size` from an ambiguous collection.*" -) - -warnings.filterwarnings( - "ignore", ".*pyprof will be removed by the end of June.*" -) - -warnings.filterwarnings( - "ignore", ".*pandas.Int64Index is deprecated.*" -) - -warnings.filterwarnings( - "ignore", ".*does not have many workers which may be a bottleneck*" -) - -warnings.filterwarnings( - "ignore", ".*Our suggested max number of worker in current system is*" -) - -os.environ["HYDRA_FULL_ERROR"] = "1" -os.environ["NUMEXPR_MAX_THREADS"] = "8" -os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:32" - -# set to medium or high -torch.set_float32_matmul_precision('medium') -import torch.multiprocessing -torch.multiprocessing.set_sharing_strategy('file_system') -torch.cuda.empty_cache() \ No newline at end of file diff --git a/personalised/code/launch/tools.py b/personalised/code/launch/tools.py deleted file mode 100644 index ce82ba7ae21a90fbe5bee1834f692ef1c98eb88d..0000000000000000000000000000000000000000 --- a/personalised/code/launch/tools.py +++ /dev/null @@ -1,7 +0,0 @@ -from omegaconf import DictConfig -import os - - -def resolve_cfg_path(cfg: DictConfig): - working_dir = os.getcwd() - cfg.working_dir = working_dir diff --git a/personalised/code/main.py b/personalised/code/main.py deleted file mode 100644 index 1e217823bc8f68699ed0b690300c02bf20614300..0000000000000000000000000000000000000000 --- a/personalised/code/main.py +++ /dev/null @@ -1,50 +0,0 @@ -import os -import sys -import hydra -import launch.prepare -from omegaconf import DictConfig -from hydra.utils import instantiate -import logging -from utils.util import set_seed - -logger = logging.getLogger(__name__) - - -@hydra.main( - version_base=None, config_path='configs', config_name='main') -def main(cfg: DictConfig): - seed = cfg.get("seed", 42) - set_seed(seed, deterministic=cfg.get("deterministic", True)) - return _main(cfg) - - -def _main(cfg): - logger.info(f"Original working directory: {hydra.utils.get_original_cwd()}") - - logger.info("Instantiating trainer") - trainer = instantiate(cfg.trainer, task=cfg.task, _recursive_=False) - logger.info("Done instantiating trainer") - - # set dataloader in trainer - logger.info("Instantiating data module") - data_module = instantiate(cfg.data, seed=cfg.get("seed", 42), _recursive_=False) - trainer.set_data_module(data_module) - - if cfg.stage == "fit": - logger.info("Model Fitting ...") - trainer.fit() - logger.info("Fitting done") - else: # test stage - logger.info("Model Testing ...") - trainer.test() - logger.info("Testing done") - - -if __name__ == '__main__': - logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' - ) - logger = logging.getLogger(__name__) - - main() \ No newline at end of file diff --git a/personalised/code/requirements.txt b/personalised/code/requirements.txt deleted file mode 100644 index a5a9e427f680e66be64ab84264b871eb5c44e1f0..0000000000000000000000000000000000000000 --- a/personalised/code/requirements.txt +++ /dev/null @@ -1,141 +0,0 @@ -absl-py -accelerate -albumentations -annotated-types -antlr4-python3-runtime -attrs -audioread -av -beartype -bibtexparser -cffi -charset-normalizer -colorama -coloredlogs -colorlog -contourpy -cycler -Cython -decorator -decord -diffusers -dlinfo -einops -et_xmlfile -face-alignment -filelock -flatbuffers -fonttools -fsspec -grpcio -h5py -huggingface-hub -humanfriendly -hydra-core -idna -imageio -imageio-ffmpeg -importlib_metadata -iopath -jax -jaxlib -Jinja2 -joblib -jsonschema -kiwisolver -lazy_loader -librosa -llvmlite -lxml -Markdown -markdown-it-py -MarkupSafe -matplotlib -mdurl -mkl-service -ml_collections -ml_dtypes -more-itertools -moviepy -msgpack -networkx -numba -numpy -omegaconf -onnx2torch-py313 -onnxruntime -openai-whisper -opencv-python-headless -openpyxl -opt_einsum -packaging -pandas -phonemizer -pillow -pipdeptree -platformdirs -plotly -pooch -portalocker -proglog -protobuf -psutil -pydantic -pydantic_core -pydub -Pygments -pyparsing -python-dateutil -python-dotenv -pytorch3d -pytz -PyYAML -rdflib -regex -requests -resampy -rich -rotary-embedding-torch -rpds-py -safetensors -samplerate -scikit-image -scikit-learn -scikit-video -scipy -segments -sentencepiece -setuptools -shortuuid -simsimd -six -sounddevice -soundfile -soxr -stringzilla -sympy -tabulate -tenacity -tensorboard -tensorboard-data-server -termcolor -threadpoolctl -tifffile -tiktoken -timm -tokenizers -torch -torchaudio -torchvision -tqdm -transformers -triton -tslearn -typing_extensions -tzdata -uritemplate -urllib3 -Werkzeug -wheel -yacs -zipp \ No newline at end of file diff --git a/personalised/code/trainer/__init__.py b/personalised/code/trainer/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/personalised/code/trainer/_archive.py b/personalised/code/trainer/_archive.py deleted file mode 100644 index 94c5b67304d1d9d78d247e0cb8b10dfe249423ee..0000000000000000000000000000000000000000 --- a/personalised/code/trainer/_archive.py +++ /dev/null @@ -1,255 +0,0 @@ -# import math -# import os -# from pathlib import Path -# import random -# import hydra -# from omegaconf import DictConfig -# import torch -# import torch.nn as nn -# from framework.utils.util import from_pretrained_checkpoint -# from utils.util import get_tensorboard_path, AverageMeter, get_lr -# from torch import optim -# from torch.nn.utils import clip_grad_norm_ -# from tqdm import tqdm -# from hydra.utils import instantiate -# from torch.utils.tensorboard import SummaryWriter -# import logging -# -# logger = logging.getLogger(__name__) -# -# -# class Trainer: -# def __init__(self, -# resumed_training: bool = False, -# model: DictConfig = None, -# criterion: DictConfig = None, -# **kwargs): -# super().__init__() -# self.resumed_training = resumed_training -# self.model_cfg = model -# self.criterion_cfg = criterion -# -# if torch.cuda.device_count() > 0: -# device = torch.device('cuda:0') -# else: -# device = torch.device('cpu') -# self.device = device -# -# self.clip_length = kwargs.pop("clip_length") -# self.window_size = kwargs.pop("window_size") * kwargs.pop("s_ratio") -# self.start_epoch = kwargs.pop("start_epoch") -# self.epochs = kwargs.pop("epochs") -# self.tb_dir = kwargs.pop("tb_dir") -# self.val_period = kwargs.pop("val_period") -# self.lr = kwargs.pop("lr") -# self.weight_decay = kwargs.pop("weight_decay") -# self.beta = kwargs.pop("beta") -# self.kwargs = kwargs -# -# def set_data_module(self, data_module): -# self.data_module = data_module -# -# def get_ckpt_path(self, model, runid="current_runid", epoch=None, best=False, last=False): -# ckpt_dir = Path(hydra.utils.to_absolute_path(self.kwargs.get("ckpt_dir"))) -# run_id = Path(self.kwargs.get(runid)) -# ckpt_dir = str(ckpt_dir / run_id / model.get_model_name()) -# os.makedirs(ckpt_dir, exist_ok=True) -# -# ckpt_path = None -# if epoch is not None: -# ckpt_path = os.path.join(ckpt_dir, f"checkpoint_{epoch}.pth") -# if best: -# ckpt_path = os.path.join(ckpt_dir, "checkpoint_best.pth") -# if last: -# ckpt_path = os.path.join(ckpt_dir, "checkpoint_last.pth") -# assert ckpt_path is not None, "No checkpoint path is provided." -# return ckpt_path -# -# def data_resample(self, emotion_clips, params_clips, seq_lengths): -# emotion_clip_list = [] -# params_clip_list = [] -# for emotion_clip, params_clip, seq_length in zip(emotion_clips, params_clips, seq_lengths): -# emotion_clip = emotion_clip[:seq_length] -# params_clip = params_clip[:seq_length] -# -# if seq_length < self.clip_length: -# cp = random.randint(0, seq_length - self.window_size) if seq_length > self.window_size else 0 -# else: -# cp = random.randint(0, self.clip_length - self.window_size) -# -# if seq_length < self.window_size: -# emotion_clip = torch.cat((emotion_clip, torch.zeros(size=(self.window_size - seq_length, -# emotion_clip.shape[-1]))), dim=0) -# params_clip = torch.cat((params_clip, torch.zeros(size=(self.window_size - seq_length, -# params_clip.shape[-1]))), dim=0) -# -# emotion_clip = emotion_clip[cp:cp + self.window_size] -# emotion_clip_list.append(emotion_clip) -# params_clip = params_clip[cp:cp + self.window_size] -# params_clip_list.append(params_clip) -# -# emotion_clips = torch.stack(emotion_clip_list, dim=0) # (bs, w, d) -# params_clips = torch.stack(params_clip_list, dim=0) -# -# return emotion_clips, params_clips -# -# def fit(self): -# stage = "fit" -# -# logger.info("Loading data module") -# self.train_loader, self.val_loader = self.data_module.get_dataloader(stage=stage) -# logger.info("Data module loaded") -# -# logger.info("Loading criterion") -# self.criterion = instantiate(self.criterion_cfg) -# logger.info("Criterion loaded") -# -# logger.info("Loading writer") -# self.writer = SummaryWriter(self.tb_dir) -# logger.info(f"Writer loaded: {self.tb_dir}") -# -# self.main_autoencoder() -# -# def compute_metrics(self, original, reconstructed): -# """ -# Compute evaluation metrics between original and reconstructed emotion data -# """ -# # Example metrics: MSE, MAE -# total_mse = 0 -# total_mae = 0 -# count = 0 -# -# for orig, recon in zip(original, reconstructed): -# # Make sure they have the same length -# min_len = min(orig.size(0), recon.size(0)) -# orig = orig[:min_len] -# recon = recon[:min_len] -# -# mse = torch.mean((orig - recon) ** 2) -# mae = torch.mean(torch.abs(orig - recon)) -# -# total_mse += mse.item() -# total_mae += mae.item() -# count += 1 -# -# avg_mse = total_mse / count if count > 0 else 0 -# avg_mae = total_mae / count if count > 0 else 0 -# -# logger.info(f"Test MSE: {avg_mse:.5f}") -# logger.info(f"Test MAE: {avg_mae:.5f}") -# -# def main_autoencoder(self): -# model = instantiate(self.model_cfg, -# _recursive_=False) -# model.to(self.device) -# -# optimizer = torch.optim.AdamW(params=model.parameters(), lr=self.lr, -# weight_decay=self.weight_decay, betas=self.beta) -# if self.resumed_training: -# checkpoint_path = self.get_ckpt_path(model.autoencoder, runid="resume_runid", best=True) -# from_pretrained_checkpoint(checkpoint_path, optimizer, self.device) -# -# best_recon_loss = float('inf') -# for epoch in range(self.start_epoch, self.epochs): -# train_loss, rec_loss, kld_loss, coeff_loss = self.train( -# model, self.train_loader, optimizer, self.criterion, epoch, self.writer, self.device -# ) -# logging.info( -# "Epoch: {} train_whole_loss: {:.5f} train_rec_loss: {:.5f} train_kld_loss: {:.5f} train_coeff_loss: {:.5f}" -# .format(epoch + 1, train_loss, rec_loss, kld_loss, coeff_loss)) -# -# if (epoch + 1) % self.val_period == 0: -# val_loss, rec_loss, kld_loss, coeff_loss = ( -# self.val(model, self.val_loader, self.criterion, self.device)) -# checkpoint_path = self.get_ckpt_path(model, epoch=(epoch + 1)) -# model.save_ckpt(checkpoint_path, optimizer) -# -# logging.info( -# "Epoch: {} val_whole_loss: {:.5f} val_rec_loss: {:.5f} val_kld_loss: {:.5f} val_coeff_loss: {:.5f}" -# .format(epoch + 1, val_loss, rec_loss, kld_loss, coeff_loss)) -# -# if val_loss < best_recon_loss: -# best_recon_loss = val_loss -# logging.info( -# f"New best reconstruction loss ({best_recon_loss:.5f}) at epoch {epoch + 1}, saving checkpoint." -# ) -# checkpoint_path = self.get_ckpt_path(model, best=True) -# model.save_ckpt(checkpoint_path, optimizer) -# -# checkpoint_path = self.get_ckpt_path(model, last=True) -# model.save_ckpt(checkpoint_path, optimizer) -# -# def train(self, model, data_loader, optimizer, criterion, epoch, writer, device): -# whole_losses = AverageMeter() -# rec_losses = AverageMeter() -# kld_losses = AverageMeter() -# div_losses = AverageMeter() -# -# model.train() -# for batch_idx, (_, _, -# emotion_clips, -# params_clips, -# _, _, _, _, -# seq_lengths) in enumerate(tqdm(data_loader)): -# -# emotion_data, params_data = self.data_resample(emotion_clips, params_clips, seq_lengths) -# emotion_data = emotion_data.to(device) -# params_data = params_data.to(device) -# batch_size = emotion_data.shape[0] -# -# outputs = model(emotion=emotion_data, _3dmm=params_data) -# loss_output = criterion(**outputs) -# loss, rec_loss, kld_loss, div_loss = ( -# loss_output["loss"], loss_output["mse"], loss_output["kld"], loss_output["coeff"] -# ) -# -# # Log metrics -# iteration = batch_idx + len(data_loader) * epoch -# if writer is not None: -# writer.add_scalar("Train/rec_loss", rec_loss.data.item(), iteration) -# writer.add_scalar("Train/kld_loss", kld_loss.data.item(), iteration) -# writer.add_scalar("Train/div_loss", div_loss.data.item(), iteration) -# -# whole_losses.update(loss.data.item(), batch_size) -# rec_losses.update(rec_loss.data.item(), batch_size) -# kld_losses.update(kld_loss.data.item(), batch_size) -# div_losses.update(div_loss.data.item(), batch_size) -# -# # Backward pass -# optimizer.zero_grad() -# loss.backward() -# optimizer.step() -# -# return whole_losses.avg, rec_losses.avg, kld_losses.avg, div_losses.avg -# -# def val(self, model, data_loader, criterion, device): -# whole_losses = AverageMeter() -# rec_losses = AverageMeter() -# kld_losses = AverageMeter() -# div_losses = AverageMeter() -# -# model.eval() -# for batch_idx, (_, _, -# emotion_clips, -# params_clips, -# _, _, _, _, -# seq_lengths) in enumerate(tqdm(data_loader)): -# emotion_data, params_data = self.data_resample(emotion_clips, params_clips, seq_lengths) -# emotion_data = emotion_data.to(device) -# params_data = params_data.to(device) -# batch_size = emotion_data.shape[0] -# -# with torch.no_grad(): -# outputs = model(emotion=emotion_data, _3dmm=params_data) -# loss_output = criterion(**outputs) -# loss, rec_loss, kld_loss, div_loss = ( -# loss_output["loss"], loss_output["mse"], loss_output["kld"], loss_output["coeff"] -# ) -# # print(f"batch {batch_idx}, loss = {loss.data.item():.6f}") -# -# whole_losses.update(loss.data.item(), batch_size) -# rec_losses.update(rec_loss.data.item(), batch_size) -# kld_losses.update(kld_loss.data.item(), batch_size) -# div_losses.update(div_loss.data.item(), batch_size) -# -# return whole_losses.avg, rec_losses.avg, kld_losses.avg, div_losses.avg diff --git a/personalised/code/trainer/emotion_autoencoder.py b/personalised/code/trainer/emotion_autoencoder.py deleted file mode 100644 index 5063b16293116e530b64f26c04f1328863d41d6e..0000000000000000000000000000000000000000 --- a/personalised/code/trainer/emotion_autoencoder.py +++ /dev/null @@ -1,338 +0,0 @@ -import os -from pathlib import Path -import hydra -import numpy as np -from omegaconf import DictConfig -import torch -import torch.nn.functional as F -from sklearn.metrics import accuracy_score, precision_recall_fscore_support, f1_score -from framework.utils.util import from_pretrained_checkpoint -from utils.util import AverageMeter, get_lr -from tqdm import tqdm -from hydra.utils import instantiate -from torch.utils.tensorboard import SummaryWriter -import logging - -logger = logging.getLogger(__name__) - - -class Trainer: - def __init__(self, - resumed_training: bool = False, - model: DictConfig = None, - criterion: DictConfig = None, - **kwargs): - super().__init__() - self.resumed_training = resumed_training - self.model_cfg = model - self.criterion_cfg = criterion - - if torch.cuda.device_count() > 0: - device = torch.device('cuda:0') - else: - device = torch.device('cpu') - self.device = device - - self.clip_length = kwargs.pop("clip_length") - self.start_epoch = kwargs.pop("start_epoch") - self.epochs = kwargs.pop("epochs") - self.tb_dir = kwargs.pop("tb_dir") - self.val_period = kwargs.pop("val_period") - self.lr = kwargs.pop("lr") - self.weight_decay = kwargs.pop("weight_decay") - self.beta = kwargs.pop("beta") - self.kwargs = kwargs - - def set_data_module(self, data_module): - self.data_module = data_module - - def get_ckpt_path(self, model, runid="current_runid", epoch=None, best=False, last=False): - ckpt_dir = Path(hydra.utils.to_absolute_path(self.kwargs.get("ckpt_dir"))) - run_id = Path(self.kwargs.get(runid)) - ckpt_dir = str(ckpt_dir / run_id / model.get_model_name()) - os.makedirs(ckpt_dir, exist_ok=True) - - ckpt_path = None - if epoch is not None: - ckpt_path = os.path.join(ckpt_dir, f"checkpoint_{epoch}.pth") - if best: - ckpt_path = os.path.join(ckpt_dir, "checkpoint_best.pth") - if last: - ckpt_path = os.path.join(ckpt_dir, "checkpoint_last.pth") - assert ckpt_path is not None, "No checkpoint path is provided." - return ckpt_path - - def compute_metrics(self, original, reconstructed): - """ - Compute evaluation metrics between original and reconstructed emotion data - """ - # Example metrics: MSE, MAE - mse = np.mean((original - reconstructed) ** 2) - mae = np.mean(np.abs(original - reconstructed)) - - return { - "mse": mse, - "mae": mae, - } - - def eval_au_binary(self, pred_list, tgt_list): - assert len(pred_list) == len(tgt_list) - accs, precs, recs, f1s = [], [], [], [] - - for preds, tgts in zip(pred_list, tgt_list): - preds_bin = preds.astype(int) - - print("number of positive samples: ", np.sum(tgts)) - print("number of negative samples: ", np.sum(1-tgts)) - print("number of positive samples predicted: ", np.sum(preds_bin)) - print("number of negative samples predicted: ", np.sum(1-preds_bin)) - - accs.append(accuracy_score(tgts, preds_bin)) - p, r, f1, _ = precision_recall_fscore_support( - tgts, - preds_bin, - average='binary', - zero_division=0 - ) - precs.append(p) - recs.append(r) - f1s.append(f1) - - results = { - 'accuracy': float(np.mean(accs)), - 'precision': float(np.mean(precs)), - 'recall': float(np.mean(recs)), - 'f1': float(np.mean(f1s)), - } - # 'micro_f1': micro_f1, - # 'macro_f1': macro_f1, - return results - - def test(self): - stage = "test" - device = self.device - logger.info("Loading test data module") - test_loader = self.data_module.get_dataloader( - stage=stage, collate_fn='none') - logger.info("Test data module loaded") - - model = instantiate(self.model_cfg, - _recursive_=False) - ckpt_path = self.get_ckpt_path(model, runid="resume_runid", best=True) - from_pretrained_checkpoint(ckpt_path, model, device) - model.eval() - - au_predictions_all = [] - va_predictions_all = [] - emotion_predictions_all = [] - emotion_targets_all = [] - au_targets_all = [] - va_targets_all = [] - - logger.info("Model testing started ...") - for batch_idx, (input_emotion_clips, e_start_indices, e_end_indices, - output_emotion_clips, d_start_indices, d_end_indices) \ - in enumerate(tqdm(test_loader)): - - (input_emotion_clips, e_start_indices, e_end_indices, d_start_indices, d_end_indices) = \ - (input_emotion_clips.to(device), e_start_indices.to(device), e_end_indices.to(device), - d_start_indices.to(device), d_end_indices.to(device)) - - with torch.no_grad(): - out, _, dist, mask = ( - model(input_emotion_clips, e_start_indices, e_end_indices, d_start_indices, d_end_indices, - reparameterization="deterministic") - ) - - lengths = (d_end_indices - d_start_indices + 1).detach().cpu().numpy() - - au_logits, va_logits, emotion_logits = out - au_predictions = (F.sigmoid(au_logits) >= 0.5).float() - au_predictions = (au_predictions * mask).detach().cpu().numpy() - au_prediction_list = np.array([]) - for au_pred, length in zip(au_predictions, lengths): - au_prediction_list = np.concatenate([au_prediction_list, au_pred[:length].reshape(-1)], axis=0) - au_predictions_all.append(au_prediction_list) - - va_predictions_all.append((va_logits * mask).detach().cpu().numpy()) - emotion_predictions_all.append((F.softmax(emotion_logits, dim=-1) * mask).detach().cpu().numpy()) - - au_targets, va_targets, emotion_targets = \ - (output_emotion_clips[:, :, :15].numpy(), - output_emotion_clips[:, :, 15:17].numpy(), - output_emotion_clips[:, :, 17:].numpy()) - - au_target_list = np.array([]) - for au_target, length in zip(au_targets, lengths): - au_target_list = np.concatenate([au_target_list, au_target[:length].reshape(-1)], axis=0) - au_targets_all.append(au_target_list) - - va_targets_all.append(va_targets) - emotion_targets_all.append(emotion_targets) - - va_predictions_all = np.concatenate(va_predictions_all, axis=0) - va_targets_all = np.concatenate(va_targets_all, axis=0) - emotion_predictions_all = np.concatenate(emotion_predictions_all, axis=0) - emotion_targets_all = np.concatenate(emotion_targets_all, axis=0) - - au_results = self.eval_au_binary(au_predictions_all, au_targets_all) - logger.info(f"AU results: {au_results}") - va_mse_results = self.compute_metrics(va_targets_all, va_predictions_all) - logger.info(f"VA results: {va_mse_results}") - emotion_mse_results = self.compute_metrics(emotion_targets_all, emotion_predictions_all) - logger.info(f"Emotion results: {emotion_mse_results}") - - def fit(self): - stage = "fit" - - logger.info("Loading data module") - self.train_loader, self.val_loader = self.data_module.get_dataloader( - stage=stage, collate_fn='none') - logger.info("Data module loaded") - - logger.info("Loading criterion") - self.criterion = instantiate(self.criterion_cfg) - logger.info("Criterion loaded") - - logger.info("Loading writer") - self.writer = SummaryWriter(self.tb_dir) - logger.info(f"Writer loaded: {self.tb_dir}") - self.main_autoencoder() - - def main_autoencoder(self): - model = instantiate(self.model_cfg, - _recursive_=False) - model.to(self.device) - - # Load optimizer - optimizer = torch.optim.AdamW(params=model.parameters(), lr=self.lr, - weight_decay=self.weight_decay, betas=self.beta) - - if self.resumed_training: - checkpoint_path = self.get_ckpt_path(model, runid="resume_runid", last=True) - from_pretrained_checkpoint(checkpoint_path, optimizer, self.device) - best_val_loss, self.start_epoch = from_pretrained_checkpoint(checkpoint_path, model, self.device) - logger.info(f"Resume training from epoch {self.start_epoch}") - else: - best_val_loss = float('inf') - print(f"Best validation loss: {best_val_loss}") - - for epoch in range(self.start_epoch, self.epochs): - train_loss, kld_loss, au_loss, va_loss, em_loss = self.train_autoencoder( - model, self.train_loader, optimizer, - self.criterion, epoch, self.writer, self.device - ) - logging.info(f"Epoch: {epoch + 1} train_loss: {train_loss:.5f} kld_loss: {kld_loss:.5f} " - f"au_loss: {au_loss:.5f} va_loss: {va_loss:.5f} em_loss: {em_loss:.5f}") - - if (epoch + 1) % self.val_period == 0: - val_loss, kld_loss, au_loss, va_loss, em_loss = self.val_autoencoder( - model, self.val_loader, self.criterion, self.device - ) - logging.info(f"Epoch: {epoch + 1} val_loss: {val_loss:.5f} kld_loss: {kld_loss:.5f} " - f"au_loss: {au_loss:.5f} va_loss: {va_loss:.5f} em_loss: {em_loss:.5f}") - - checkpoint = { - 'epoch': epoch + 1, - 'best_loss': best_val_loss, - 'state_dict': model.state_dict(), - 'optimizer': optimizer.state_dict(), - } - ckpt_path = self.get_ckpt_path(model, epoch=(epoch + 1)) - torch.save(checkpoint, ckpt_path) - ckpt_path = self.get_ckpt_path(model, last=True) - torch.save(checkpoint, ckpt_path) - - if val_loss < best_val_loss: - best_val_loss = val_loss - logging.info( - f"New best loss ({best_val_loss:.5f}) at epoch {epoch + 1}, saving checkpoint." - ) - ckpt_path = self.get_ckpt_path(model, best=True) - torch.save(checkpoint, ckpt_path) - - def train_autoencoder(self, model, data_loader, optimizer, - criterion, epoch, writer, device): - whole_losses = AverageMeter() - kld_losses = AverageMeter() - au_losses = AverageMeter() - va_losses = AverageMeter() - em_losses = AverageMeter() - - model.train() - for batch_idx, (input_emotion_clips, e_start_indices, e_end_indices, - output_emotion_clips, d_start_indices, d_end_indices) \ - in enumerate(tqdm(data_loader)): - - (input_emotion_clips, output_emotion_clips, e_start_indices, - e_end_indices, d_start_indices, d_end_indices) = \ - (input_emotion_clips.to(device), output_emotion_clips.to(device), e_start_indices.to(device), - e_end_indices.to(device), d_start_indices.to(device), d_end_indices.to(device)) - batch_size = input_emotion_clips.shape[0] - - out, _, dist, mask = ( - model(input_emotion_clips, e_start_indices, e_end_indices, d_start_indices, d_end_indices) - ) - # out: ([B x L x 15], [B x L x 2], [B x L x 8]) - - loss, kld_loss, au_loss, va_loss, em_loss = \ - criterion(predictions=out, targets=output_emotion_clips, distribution=dist, mask=mask) - - # Log metrics - iteration = batch_idx + len(data_loader) * epoch - if writer is not None: - writer.add_scalar("Train/whole_loss", loss.data.item(), iteration) - writer.add_scalar("Train/kld_loss", kld_loss.data.item(), iteration) - writer.add_scalar("Train/au_loss", au_loss.data.item(), iteration) - writer.add_scalar("Train/va_loss", va_loss.data.item(), iteration) - writer.add_scalar("Train/em_loss", em_loss.data.item(), iteration) - - whole_losses.update(loss.data.item(), batch_size) - kld_losses.update(kld_loss.data.item(), batch_size) - au_losses.update(au_loss.data.item(), batch_size) - va_losses.update(va_loss.data.item(), batch_size) - em_losses.update(em_loss.data.item(), batch_size) - - # Backward pass - optimizer.zero_grad() - loss.backward() - optimizer.step() - - # Get learning rate - lr = get_lr(optimizer=optimizer) - if writer is not None: - writer.add_scalar("Train/lr", lr, epoch) - - return whole_losses.avg, kld_losses.avg, au_losses.avg, va_losses.avg, em_losses.avg - - def val_autoencoder(self, model, val_loader, criterion, device): - whole_losses = AverageMeter() - kld_losses = AverageMeter() - au_losses = AverageMeter() - va_losses = AverageMeter() - em_losses = AverageMeter() - - model.eval() - for batch_idx, (input_emotion_clips, e_start_indices, e_end_indices, - output_emotion_clips, d_start_indices, d_end_indices) \ - in enumerate(tqdm(val_loader)): - (input_emotion_clips, output_emotion_clips, e_start_indices, - e_end_indices, d_start_indices, d_end_indices) = \ - (input_emotion_clips.to(device), output_emotion_clips.to(device), e_start_indices.to(device), - e_end_indices.to(device), d_start_indices.to(device), d_end_indices.to(device)) - batch_size = input_emotion_clips.shape[0] - - with torch.no_grad(): - out, _, dist, mask = ( - model(input_emotion_clips, e_start_indices, e_end_indices, d_start_indices, d_end_indices) - ) - loss, kld_loss, au_loss, va_loss, em_loss = \ - criterion(predictions=out, targets=output_emotion_clips, distribution=dist, mask=mask) - - whole_losses.update(loss.data.item(), batch_size) - kld_losses.update(kld_loss.data.item(), batch_size) - au_losses.update(au_loss.data.item(), batch_size) - va_losses.update(va_loss.data.item(), batch_size) - em_losses.update(em_loss.data.item(), batch_size) - - return whole_losses.avg, kld_losses.avg, au_losses.avg, va_losses.avg, em_losses.avg \ No newline at end of file diff --git a/personalised/code/trainer/g2p_delta.py b/personalised/code/trainer/g2p_delta.py deleted file mode 100644 index 68161d747b3e4a12e172bb207b3b7fdb0902a7ce..0000000000000000000000000000000000000000 --- a/personalised/code/trainer/g2p_delta.py +++ /dev/null @@ -1,408 +0,0 @@ -"""Training/evaluation glue for G2P-Delta on the personalised MARS loader.""" - -import torch -import torch.nn.functional as F -from torch import optim -from tqdm import tqdm - -from framework.g2p_delta import G2PDeltaModel -from trainer.perfrdiff_rewrite_weight import Trainer as _PersonalisedTrainer -from utils.util import AverageMeter - - -def _ccc_loss(prediction, target, eps=1.0e-6): - prediction = prediction.reshape(-1, prediction.shape[-2], prediction.shape[-1]) - target = target.reshape(-1, target.shape[-2], target.shape[-1]) - pred_mean = prediction.mean(dim=1) - target_mean = target.mean(dim=1) - pred_centered = prediction - pred_mean.unsqueeze(1) - target_centered = target - target_mean.unsqueeze(1) - covariance = (pred_centered * target_centered).mean(dim=1) - pred_var = pred_centered.square().mean(dim=1) - target_var = target_centered.square().mean(dim=1) - ccc = 2.0 * covariance / ( - pred_var + target_var + (pred_mean - target_mean).square() + eps - ) - return 1.0 - ccc.mean() - - -def _dynamics_loss(prediction, target): - prediction = prediction.reshape(-1, prediction.shape[-2], prediction.shape[-1]) - target = target.reshape(-1, target.shape[-2], target.shape[-1]) - pred_velocity = prediction[:, 1:] - prediction[:, :-1] - target_velocity = target[:, 1:] - target[:, :-1] - velocity = F.mse_loss(pred_velocity, target_velocity) - if prediction.shape[1] < 3: - return velocity - pred_accel = pred_velocity[:, 1:] - pred_velocity[:, :-1] - target_accel = target_velocity[:, 1:] - target_velocity[:, :-1] - return velocity + 0.5 * F.mse_loss(pred_accel, target_accel) - - -def _coarse_loss(output_decoder): - logits = output_decoder.get("coarse_logits") - if logits is None: - return output_decoder["prediction_emotion"].new_tensor(0.0) - target = output_decoder["target_emotion"][..., 17:25] - denom = target.sum(dim=-1, keepdim=True) - valid = (denom > 0.5).to(logits.dtype) - probabilities = target / denom.clamp_min(1.0e-6) - cross_entropy = -(probabilities * F.log_softmax(logits, dim=-1)).sum( - dim=-1, keepdim=True - ) - return (cross_entropy * valid).sum() / valid.sum().clamp_min(1.0) - - -class Trainer(_PersonalisedTrainer): - def _build_model(self, stage): - diffusion = self._build_diffusion(stage) - model = G2PDeltaModel(self._rewrite_cfg(), diffusion) - model.to(self.device) - return model - - def _build_optimizer(self, model): - args = self.main_model_cfg.optimizer_hypernet.args - adapter_params = list(model.modifier_parameters(include_eeg_head=False)) - groups = [{"params": adapter_params, "lr": float(args.lr)}] - if self.train_eeg: - model.set_eeg_head_requires_grad(True) - groups.append( - { - "params": list(model.eeg_head().parameters()), - "lr": float(self.main_model_cfg.args.get("eeg_lr", 1.0e-5)), - } - ) - return optim.AdamW(groups, weight_decay=float(args.weight_decay)) - - def _batch_loss( - self, - model, - criterion, - speaker_audio, - speaker_emotion, - speaker_3dmm, - listener_emotion, - past_listener_emotion, - motion_length, - personal_3dmm, - listener_personality, - listener_eeg, - listener_eeg_mask, - counterfactual, - ): - input_dict = { - "speaker_audio_input": speaker_audio, - "speaker_emotion_input": speaker_emotion, - "speaker_3dmm_input": speaker_3dmm, - "listener_emotion_input": listener_emotion, - "past_listener_emotion": past_listener_emotion, - "motion_length": motion_length, - "listener_eeg_input": listener_eeg, - "listener_eeg_mask": listener_eeg_mask, - } - personal = personal_3dmm if personal_3dmm.numel() > 0 else None - cpu_rng_before = torch.random.get_rng_state() - cuda_rng_before = ( - torch.cuda.get_rng_state(self.device) - if self.device.type == "cuda" - else None - ) - outputs, regular = model( - x=input_dict, p=personal, personality=listener_personality - ) - output_prior, output_decoder = self._split_outputs(outputs) - losses = criterion(output_prior, output_decoder) - args = self.main_model_cfg.args - coarse = _coarse_loss(output_decoder) - ccc = _ccc_loss( - output_decoder["prediction_emotion"], output_decoder["target_emotion"] - ) - dynamics = _dynamics_loss( - output_decoder["prediction_emotion"], output_decoder["target_emotion"] - ) - total = ( - losses["loss"] - + float(args.get("coarse_weight", 0.5)) * coarse - + float(args.get("ccc_weight", 0.1)) * ccc - + float(args.get("dynamics_weight", 0.02)) * dynamics - + regular - ) - - cf_loss = total.new_tensor(0.0) - cf_weight = float(args.get("counterfactual_weight", 0.1)) - if counterfactual and cf_weight > 0 and speaker_audio.shape[0] > 1: - torch.random.set_rng_state(cpu_rng_before) - if cuda_rng_before is not None: - torch.cuda.set_rng_state(cuda_rng_before, self.device) - negative_personal = ( - personal.roll(1, dims=0) if personal is not None else None - ) - negative_personality = listener_personality.roll(1, dims=0) - negative_outputs, _ = model( - x=input_dict, - p=negative_personal, - personality=negative_personality, - ) - _, negative_decoder = self._split_outputs(negative_outputs) - negative_losses = criterion(output_prior, negative_decoder) - margin = float(args.get("counterfactual_margin", 0.02)) - cf_loss = F.relu( - margin + losses["decoded"] - negative_losses["decoded"] - ) - total = total + cf_weight * cf_loss - - losses["loss"] = total - losses["loss_coarse"] = coarse - losses["loss_ccc"] = ccc - losses["loss_dynamics"] = dynamics - losses["loss_counterfactual"] = cf_loss - return total, losses, regular - - def _run_epoch(self, model, data_loader, criterion, optimizer, writer, epoch, train=True): - meters = [AverageMeter() for _ in range(6)] - whole, prior, decoded, eeg, eeg_valid, regular_meter = meters - model.train(train) - if train and self.train_eeg: - model.eeg_head().train() - max_batches = int(self.trainer_cfg.get("max_train_batches", 0)) if train else int( - self.trainer_cfg.get("max_val_batches", 0) - ) - - for batch_idx, batch in enumerate(tqdm(data_loader)): - if max_batches > 0 and batch_idx >= max_batches: - break - if len(batch) != 12: - raise ValueError( - "G2P-Delta training requires the 12-item personalised batch " - "with personality and EEG tensors." - ) - ( - speaker_audio, - _, - speaker_emotion, - speaker_3dmm, - _, - listener_emotion, - _, - personal_3dmm, - listener_personality, - listener_eeg, - listener_eeg_mask, - _, - ) = batch - tensors = [ - speaker_audio, - speaker_emotion, - speaker_3dmm, - listener_emotion, - personal_3dmm, - listener_personality, - listener_eeg, - listener_eeg_mask, - ] - ( - speaker_audio, - speaker_emotion, - speaker_3dmm, - listener_emotion, - personal_3dmm, - listener_personality, - listener_eeg, - listener_eeg_mask, - ) = [tensor.to(self.device) for tensor in tensors] - ( - speaker_audio, - speaker_emotion, - speaker_3dmm, - listener_emotion, - past_listener_emotion, - motion_length, - listener_eeg, - listener_eeg_mask, - ) = self._resample_train_batch( - speaker_audio, - speaker_emotion, - speaker_3dmm, - listener_emotion, - listener_eeg=listener_eeg, - listener_eeg_mask=listener_eeg_mask, - ) - if optimizer is not None: - optimizer.zero_grad(set_to_none=True) - context = torch.enable_grad() if train else torch.no_grad() - with context: - loss, loss_dict, regular = self._batch_loss( - model, - criterion, - speaker_audio, - speaker_emotion, - speaker_3dmm, - listener_emotion, - past_listener_emotion, - motion_length.to(self.device) if motion_length is not None else None, - personal_3dmm, - listener_personality, - listener_eeg, - listener_eeg_mask, - counterfactual=train, - ) - if train: - loss.backward() - if train: - if self.trainer_cfg.clip_grad: - torch.nn.utils.clip_grad_norm_( - [p for p in model.parameters() if p.requires_grad], 1.0 - ) - optimizer.step() - - batch_size = speaker_audio.shape[0] - whole.update(loss.detach().item(), batch_size) - prior.update(loss_dict["encoded"].detach().item(), batch_size) - decoded.update(loss_dict["decoded"].detach().item(), batch_size) - eeg.update(loss_dict["loss_eeg"].detach().item(), batch_size) - eeg_valid.update(loss_dict["eeg_valid_ratio"].detach().item(), batch_size) - regular_meter.update(regular.detach().item(), batch_size) - if writer is not None: - step = batch_idx + len(data_loader) * epoch - prefix = "Train" if train else "Val" - writer.add_scalar(f"{prefix}/loss", loss.detach().item(), step) - writer.add_scalar( - f"{prefix}/loss_ccc", loss_dict["loss_ccc"].detach().item(), step - ) - writer.add_scalar( - f"{prefix}/loss_counterfactual", - loss_dict["loss_counterfactual"].detach().item(), - step, - ) - return whole.avg, prior.avg, decoded.avg, eeg.avg, eeg_valid.avg - - def _single_loss( - self, - model, - criterion, - speaker_audio, - speaker_emotion, - speaker_3dmm, - listener_emotion, - past_listener_emotion, - motion_length, - personal_3dmm, - listener_personality, - listener_eeg, - listener_eeg_mask, - idx, - ): - input_dict = { - "speaker_audio_input": speaker_audio[idx : idx + 1], - "speaker_emotion_input": speaker_emotion[idx : idx + 1], - "speaker_3dmm_input": speaker_3dmm[idx : idx + 1], - "listener_emotion_input": listener_emotion[idx : idx + 1], - "past_listener_emotion": ( - past_listener_emotion[idx : idx + 1] - if past_listener_emotion is not None - else None - ), - "motion_length": ( - motion_length[idx : idx + 1].to(self.device) - if motion_length is not None - else None - ), - "listener_eeg_input": ( - listener_eeg[idx : idx + 1] if listener_eeg is not None else None - ), - "listener_eeg_mask": ( - listener_eeg_mask[idx : idx + 1] - if listener_eeg_mask is not None - else None - ), - } - personal = personal_3dmm[idx : idx + 1] if personal_3dmm.numel() > 0 else None - personality = listener_personality[idx : idx + 1] - cpu_rng_before = torch.random.get_rng_state() - cuda_rng_before = ( - torch.cuda.get_rng_state(self.device) - if self.device.type == "cuda" - else None - ) - outputs, regular = model(x=input_dict, p=personal, personality=personality) - output_prior, output_decoder = self._split_outputs(outputs) - losses = criterion(output_prior, output_decoder) - - args = self.main_model_cfg.args - coarse = _coarse_loss(output_decoder) - ccc = _ccc_loss( - output_decoder["prediction_emotion"], output_decoder["target_emotion"] - ) - dynamics = _dynamics_loss( - output_decoder["prediction_emotion"], output_decoder["target_emotion"] - ) - total = ( - losses["loss"] - + float(args.get("coarse_weight", 0.5)) * coarse - + float(args.get("ccc_weight", 0.1)) * ccc - + float(args.get("dynamics_weight", 0.02)) * dynamics - + regular - ) - - # Same-context listener swap. RNG is restored so matched and swapped - # conditions see exactly the same diffusion noise. - counterfactual = total.new_tensor(0.0) - cf_weight = float(args.get("counterfactual_weight", 0.1)) - if ( - torch.is_grad_enabled() - and cf_weight > 0 - and speaker_audio.shape[0] > 1 - ): - negative_idx = (idx + 1) % speaker_audio.shape[0] - negative_personal = ( - personal_3dmm[negative_idx : negative_idx + 1] - if personal_3dmm.numel() > 0 - else None - ) - negative_personality = listener_personality[ - negative_idx : negative_idx + 1 - ] - torch.random.set_rng_state(cpu_rng_before) - if cuda_rng_before is not None: - torch.cuda.set_rng_state(cuda_rng_before, self.device) - negative_outputs, _ = model( - x=input_dict, - p=negative_personal, - personality=negative_personality, - ) - _, negative_decoder = self._split_outputs(negative_outputs) - negative_losses = criterion(output_prior, negative_decoder) - margin = float(args.get("counterfactual_margin", 0.02)) - counterfactual = F.relu( - margin + losses["decoded"] - negative_losses["decoded"] - ) - total = total + cf_weight * counterfactual - - losses["loss"] = total - losses["loss_coarse"] = coarse - losses["loss_ccc"] = ccc - losses["loss_dynamics"] = dynamics - losses["loss_counterfactual"] = counterfactual - return total, losses, regular - - def _apply_personalization(self, model, personal_3dmm, listener_personality): - # Short-triage control: "identity" reproduces the frozen Generic - # backbone exactly (no listener condition ever set); "matched" and - # "shuffled" both flow through the normal set_person_condition path - # and only differ in which personality/history the eval dataset - # handed us (see scripts/build_subset_eval_root.py, which builds a - # "shuffled" personality.csv variant so the swap happens at the data - # layer, not here). - mode = str(self.trainer_cfg.get("eval_condition_mode", "matched")) - if mode == "identity": - model.clear_person_condition() - return - if mode not in {"matched", "shuffled"}: - raise ValueError(f"Unknown trainer.generic.eval_condition_mode: {mode}") - personal = ( - personal_3dmm.unsqueeze(0).to(self.device) - if personal_3dmm.numel() > 0 - else None - ) - personality = listener_personality.unsqueeze(0).to(self.device) - model.set_person_condition(p=personal, personality=personality) diff --git a/personalised/code/trainer/motion_diffusion.py b/personalised/code/trainer/motion_diffusion.py deleted file mode 100644 index bad0d97fea0631f341646f977dfd8919217bebbc..0000000000000000000000000000000000000000 --- a/personalised/code/trainer/motion_diffusion.py +++ /dev/null @@ -1,1031 +0,0 @@ -import math -import os -import random -from einops import rearrange -import torch -from framework.modules.post_processor import Processor -from framework.utils.compute_metrics import compute_eeg_metrics, compute_metrics -from framework.utils.util import from_pretrained_checkpoint -from utils.util import AverageMeter, get_lr -from omegaconf import DictConfig -from tqdm import tqdm -from hydra.utils import instantiate, to_absolute_path -from torch.utils.tensorboard import SummaryWriter -import logging - -logger = logging.getLogger(__name__) - - -class Trainer: - def __init__(self, - resumed_training: bool = False, - generic: DictConfig = None, - renderer: DictConfig = None, - model: DictConfig = None, - criterion: DictConfig = None, - **kwargs): - # # current working directory: outputs/${trainer.task_name}/${data.data_name}/${run_id} - # folder: save/${trainer.task_name}/${data.data_name} # ckpt_name: checkpoint.pth - # # last ckpt directory - # ckpt_dir: ${get_last_checkpoint:${trainer.folder}} # ${trainer.run_id} - # # for example, ckpt_dir: save/motion_diffusion/react_2024/checkpoints - # resume_run_id: ${old_run_id} - - super().__init__() - self.resumed_training = resumed_training - self.renderer = renderer - self.model_cfg = model - self.criterion_cfg = criterion - - if torch.cuda.device_count() > 0: - device = torch.device('cuda:0') - else: - device = torch.device('cpu') - self.device = device - self.kwargs = kwargs - self.trainer_cfg = generic - self.optim_cfg = kwargs.pop("optim") - self.task = kwargs.get("task") - self.train_eeg_head_only = self._as_bool( - self.trainer_cfg.get("train_eeg_head_only", False) - ) - - @staticmethod - def _as_bool(value): - if isinstance(value, str): - return value.lower() in {"1", "true", "yes", "y"} - return bool(value) - - @staticmethod - def _count_parameters(parameters): - return sum(parameter.numel() for parameter in parameters) - - @staticmethod - def _resolve_checkpoint_path(path): - if path is None or str(path).strip() == "": - return None - - path = str(path) - if os.path.isabs(path): - return path - return to_absolute_path(path) - - def _load_pretrained_motion_diffusion(self, model): - decoder_checkpoint = self._resolve_checkpoint_path( - self.trainer_cfg.get("pretrained_decoder_checkpoint", "") - ) - prior_checkpoint = self._resolve_checkpoint_path( - self.trainer_cfg.get("pretrained_prior_checkpoint", "") - ) - load_prior = self._as_bool(self.trainer_cfg.get("pretrained_load_prior", True)) - - if not self.train_eeg_head_only and decoder_checkpoint is None and prior_checkpoint is None: - return - - if self.train_eeg_head_only and decoder_checkpoint is None: - raise ValueError( - "train_eeg_head_only=True requires trainer.generic.pretrained_decoder_checkpoint. " - "Use resume=false and point it to a pretrained TransformerDenoiser checkpoint." - ) - - if decoder_checkpoint is not None: - if not os.path.exists(decoder_checkpoint): - raise FileNotFoundError(f"Missing pretrained decoder checkpoint: {decoder_checkpoint}") - from_pretrained_checkpoint(decoder_checkpoint, model.diffusion_decoder.model, self.device) - logger.info(f"Loaded pretrained decoder checkpoint: {decoder_checkpoint}") - - if not load_prior or model.diffusion_prior is None: - if self.train_eeg_head_only: - model.diffusion_prior = None - return - - if prior_checkpoint is None: - if self.train_eeg_head_only: - logger.warning( - "pretrained_load_prior=True but pretrained_prior_checkpoint is empty; skip prior loading." - ) - model.diffusion_prior = None - return - if not os.path.exists(prior_checkpoint): - logger.warning(f"Missing pretrained prior checkpoint; skip prior loading: {prior_checkpoint}") - if self.train_eeg_head_only: - model.diffusion_prior = None - return - - from_pretrained_checkpoint(prior_checkpoint, model.diffusion_prior.model, self.device) - logger.info(f"Loaded pretrained prior checkpoint: {prior_checkpoint}") - - def set_data_module(self, data_module): - self.data_module = data_module - - def data_resample(self, - speaker_audio_clips, speaker_emotion_clips, speaker_3dmm_clips, - listener_video_clips, listener_emotion_clips, listener_3dmm_clips, - speaker_seq_lengths, listener_seq_lengths, - listener_eeg_clips=None, listener_eeg_masks=None): - - s_ratio = self.trainer_cfg.s_ratio - window_size = self.trainer_cfg.window_size - clip_length = self.trainer_cfg.clip_length - s_window_size = s_ratio * window_size - l_window_size = window_size - - if self.task == 'offline': - stack = lambda clips: torch.stack(clips, dim=0) - speaker_audio, speaker_emotion, speaker_3dmm = ( - stack(clips) for clips in (speaker_audio_clips, speaker_emotion_clips, speaker_3dmm_clips)) - listener_video, listener_emotion, listener_3dmm = ( - stack(clips) for clips in (listener_video_clips, listener_emotion_clips, listener_3dmm_clips)) - past_listener_emotion = past_listener_3dmm = None - seq_lengths = torch.tensor(speaker_seq_lengths).clamp(max=clip_length) - listener_eeg = listener_eeg_mask = None - # Tensor([58, 750, 632, ...]) - - elif self.task == "online": - def get_padded(clip: torch.Tensor, length: int, target_len: int) -> torch.Tensor: - clip = clip[:length] - if length < target_len: - pad_shape = (target_len - length, *clip.shape[1:]) - clip = torch.cat([clip, clip.new_zeros(pad_shape)], dim=0) - return clip - - speaker_audio, speaker_emotion, speaker_3dmm = [], [], [] - listener_video, listener_emotion, listener_3dmm = [], [], [] - past_listener_emotion, past_listener_3dmm = [], [] - listener_eeg, listener_eeg_mask = [], [] - has_eeg = listener_eeg_clips is not None and listener_eeg_masks is not None - eeg_clips = listener_eeg_clips if has_eeg else [None] * len(speaker_audio_clips) - eeg_masks = listener_eeg_masks if has_eeg else [None] * len(speaker_audio_clips) - - for (speaker_audio_clip, speaker_emotion_clip, speaker_3dmm_clip, speaker_seq_length, - listener_video_clip, listener_emotion_clip, listener_3dmm_clip, listener_seq_length, - listener_eeg_clip, listener_eeg_mask_clip) in \ - zip(speaker_audio_clips, speaker_emotion_clips, speaker_3dmm_clips, speaker_seq_lengths, - listener_video_clips, listener_emotion_clips, listener_3dmm_clips, listener_seq_lengths, - eeg_clips, eeg_masks): - seq_length = speaker_seq_length - assert speaker_seq_length == listener_seq_length, "Sequence length not equal" - - speaker_audio_clip = get_padded(speaker_audio_clip, seq_length, s_window_size) - speaker_emotion_clip = get_padded(speaker_emotion_clip, seq_length, s_window_size) - speaker_3dmm_clip = get_padded(speaker_3dmm_clip, seq_length, s_window_size) - listener_video_clip = get_padded(listener_video_clip, seq_length, s_window_size) - listener_emotion_clip = get_padded(listener_emotion_clip, seq_length, s_window_size) - listener_3dmm_clip = get_padded(listener_3dmm_clip, seq_length, s_window_size) - if has_eeg: - listener_eeg_clip = get_padded(listener_eeg_clip, seq_length, s_window_size) - listener_eeg_mask_clip = get_padded(listener_eeg_mask_clip, seq_length, s_window_size) - - if seq_length < clip_length: - cp = random.randint(0, seq_length - s_window_size) if seq_length > s_window_size else 0 - else: - cp = random.randint(0, clip_length - s_window_size) - - du = cp + s_window_size - speaker_audio_clip = speaker_audio_clip[cp:du] - speaker_emotion_clip = speaker_emotion_clip[cp:du] - speaker_3dmm_clip = speaker_3dmm_clip[cp:du] - listener_video_clip = listener_video_clip[du - l_window_size:du] - - # past = the K listener windows immediately before the target window, - # i.e. frames [du-(K+1)*lw : du-lw]; front-pad with zeros when the - # sequence does not reach that far back (K=1 -> original behaviour). - K = int(self.trainer_cfg.get("n_past_win", 1)) - pe = du - l_window_size - ps = du - (K + 1) * l_window_size - - def _past_slice(clip, _ps=ps, _pe=pe): - if _ps >= 0: - return clip[_ps:_pe] - pad = clip.new_zeros((-_ps, *clip.shape[1:])) - return torch.cat([pad, clip[0:_pe]], dim=0) - - past_listener_emotion_clip = _past_slice(listener_emotion_clip) - past_listener_3dmm_clip = _past_slice(listener_3dmm_clip) - listener_emotion_clip = listener_emotion_clip[(du - l_window_size): du] - listener_3dmm_clip = listener_3dmm_clip[(du - l_window_size): du] - if has_eeg: - listener_eeg.append(listener_eeg_clip[du - 1]) - listener_eeg_mask.append(listener_eeg_mask_clip[du - 1]) - - speaker_audio.append(speaker_audio_clip) - speaker_emotion.append(speaker_emotion_clip) - speaker_3dmm.append(speaker_3dmm_clip) - listener_video.append(listener_video_clip) - listener_emotion.append(listener_emotion_clip) - listener_3dmm.append(listener_3dmm_clip) - past_listener_emotion.append(past_listener_emotion_clip) - past_listener_3dmm.append(past_listener_3dmm_clip) - - speaker_audio = torch.stack(speaker_audio, dim=0) # (bs, s_w, d) - speaker_emotion = torch.stack(speaker_emotion, dim=0) # (bs, s_w, 25) - speaker_3dmm = torch.stack(speaker_3dmm, dim=0) # (bs, s_w, 58) - listener_video = torch.stack(listener_video, dim=0) # (bs, l_w, 3, 224, 224) - listener_emotion = torch.stack(listener_emotion, dim=0) # (bs, l_w, 25) - listener_3dmm = torch.stack(listener_3dmm, dim=0) # (bs, l_w, 58) - past_listener_emotion = torch.stack(past_listener_emotion, dim=0) # (bs, l_w, 25) - past_listener_3dmm = torch.stack(past_listener_3dmm, dim=0) # (bs, l_w, 58) - if has_eeg: - listener_eeg = torch.stack(listener_eeg, dim=0) # (bs, d_eeg) - listener_eeg_mask = torch.stack(listener_eeg_mask, dim=0) # (bs, d_eeg) - else: - listener_eeg = listener_eeg_mask = None - seq_lengths = None - else: - raise ValueError("Unknown task type") - - return (speaker_audio, speaker_emotion, speaker_3dmm, listener_video, listener_emotion, - listener_3dmm, past_listener_emotion, past_listener_3dmm, seq_lengths, - listener_eeg, listener_eeg_mask) - - def fit(self): - """ - # relative directory - root_dir = save/${trainer.task_name}/${data.data_name}/${folder_name} - # absolute directory - saving_dir = Path(hydra.utils.to_absolute_path(root_dir)) - # get saving path - saving_path = str(saving_dir / ...) - """ - - self.start_epoch = self.trainer_cfg.start_epoch - self.epochs = self.trainer_cfg.epochs - self.tb_dir = self.trainer_cfg.tb_dir - self.clip_grad = self.trainer_cfg.clip_grad - self.val_period = self.trainer_cfg.val_period - stage = "fit" - - logger.info("Loading data module") - self.train_loader, self.val_loader = self.data_module.get_dataloader(stage=stage) - logger.info("Data module loaded") - - logger.info("Loading criterion") - self.criterion = instantiate(self.criterion_cfg) - logger.info("Criterion loaded") - - logger.info("Loading writer") - self.writer = SummaryWriter(self.tb_dir) - logger.info(f"Writer loaded: {self.tb_dir}") - self.main_diffusion(stage) - - def main_diffusion(self, stage): - if self.train_eeg_head_only and self.resumed_training: - raise ValueError( - "train_eeg_head_only=True should be launched with resume=false, " - "so the optimizer and old EEG head checkpoint are not restored." - ) - - model = instantiate(self.model_cfg.diff_model, - stage=stage, - resumed_training=self.resumed_training, - latent_embedder=self.model_cfg.latent_embedder \ - if hasattr(self.model_cfg, "latent_embedder") else None, - audio_encoder=self.model_cfg.audio_encoder \ - if hasattr(self.model_cfg, "audio_encoder") else None, - **self.kwargs, - _recursive_=False) - model.to(self.device) - self._load_pretrained_motion_diffusion(model) - - optimizer_params = model.parameters() - if self.train_eeg_head_only: - model.freeze_except_eeg_head() - trainable_params = [parameter for parameter in model.parameters() if parameter.requires_grad] - trainable_names = [name for name, parameter in model.named_parameters() if parameter.requires_grad] - frozen_count = self._count_parameters( - parameter for parameter in model.parameters() if not parameter.requires_grad - ) - trainable_count = self._count_parameters(trainable_params) - if len(trainable_params) == 0: - raise RuntimeError("No trainable parameters found for EEG head-only training.") - optimizer_params = trainable_params - logger.info( - "EEG head-only training enabled. " - f"Trainable parameters: {trainable_count}; frozen parameters: {frozen_count}" - ) - logger.info(f"Trainable parameter tensors: {trainable_names}") - print( - "EEG head-only training enabled. " - f"Trainable parameters: {trainable_count}; frozen parameters: {frozen_count}" - ) - print(f"Trainable parameter tensors: {trainable_names}") - - # load optimizer - optimizer = instantiate(self.optim_cfg, lr=self.trainer_cfg.lr, params=optimizer_params) - if self.resumed_training: - checkpoint_path = model.get_ckpt_path(model.diffusion_decoder.model, runid="resume_runid", last=True) - best_validation_loss, self.start_epoch = ( - from_pretrained_checkpoint(checkpoint_path, optimizer, self.device) - ) - logger.info(f"Resume training from epoch {self.start_epoch}") - else: - best_validation_loss = float('inf') - print(f"Best validation loss: {best_validation_loss}") - - # load scheduler - scheduler = instantiate(self.kwargs.pop("scheduler"), optimizer, len(self.train_loader)) - selected_loss_name = "loss_eeg" if self.train_eeg_head_only else "diff_loss" - - for epoch in range(self.start_epoch, self.epochs): - diffusion_loss, prior_loss, au_rec_loss, va_rec_loss, em_rec_loss, eeg_rec_loss, eeg_valid_ratio = ( - self.train_diffusion(model, self.train_loader, optimizer, scheduler, - self.criterion, epoch, self.writer, self.device)) - logging.info(f"Epoch: {epoch + 1} train_{selected_loss_name}: {diffusion_loss:.5f} " - f"prior_loss: {prior_loss:.5f} au_rec_loss: {au_rec_loss:.5f}" - f" va_rec_loss: {va_rec_loss:.5f} em_rec_loss: {em_rec_loss:.5f}" - f" eeg_rec_loss: {eeg_rec_loss:.5f} eeg_valid_ratio: {eeg_valid_ratio:.5f}") - # epoch-aligned train curve (Train/loss above is per-iteration; this is per-epoch - # so it can be overlaid with Epoch/val_loss to read convergence/overfit). - if self.writer is not None: - self.writer.add_scalar("Epoch/train_loss", diffusion_loss, epoch + 1) - - if (epoch + 1) % self.val_period == 0: - diffusion_loss, prior_loss, au_rec_loss, va_rec_loss, em_rec_loss, eeg_rec_loss, eeg_valid_ratio = ( - self.val_diffusion(model, self.val_loader, self.criterion, self.device)) - logging.info(f"Epoch: {epoch + 1} val_{selected_loss_name}: {diffusion_loss:.5f} " - f"prior_loss: {prior_loss:.5f} au_rec_loss: {au_rec_loss:.5f}" - f" va_rec_loss: {va_rec_loss:.5f} em_rec_loss: {em_rec_loss:.5f}" - f" eeg_rec_loss: {eeg_rec_loss:.5f} eeg_valid_ratio: {eeg_valid_ratio:.5f}") - if self.writer is not None: - self.writer.add_scalar("Epoch/val_loss", diffusion_loss, epoch + 1) - self.writer.add_scalar("Epoch/val_prior_loss", prior_loss, epoch + 1) - self.writer.add_scalar("Epoch/val_au_rec_loss", au_rec_loss, epoch + 1) - self.writer.add_scalar("Epoch/val_va_rec_loss", va_rec_loss, epoch + 1) - self.writer.add_scalar("Epoch/val_em_rec_loss", em_rec_loss, epoch + 1) - - if diffusion_loss < best_validation_loss: - best_validation_loss = diffusion_loss - logging.info( - f"New best {selected_loss_name} ({best_validation_loss:.5f}) at epoch {epoch + 1}, " - f"saving checkpoint" - ) - model.save_ckpt(optimizer, best=True, epoch=(epoch+1), best_loss=best_validation_loss) - - model.save_ckpt(optimizer, epoch=(epoch + 1), best_loss=best_validation_loss) - model.save_ckpt(optimizer, last=True, epoch=(epoch+1), best_loss=best_validation_loss) - - def train_diffusion(self, model, data_loader, optimizer, scheduler, - criterion, epoch, writer, device): - whole_losses = AverageMeter() - prior_losses = AverageMeter() - au_rec_losses = AverageMeter() - va_rec_losses = AverageMeter() - em_rec_losses = AverageMeter() - eeg_rec_losses = AverageMeter() - eeg_valid_ratios = AverageMeter() - - if self.train_eeg_head_only: - model.set_eeg_head_train_mode() - else: - model.train() - for batch_idx, batch in enumerate(tqdm(data_loader)): - ( - speaker_audio_clip, - speaker_video_clip, - speaker_emotion_clip, - speaker_3dmm_clip, - listener_video_clip, - listener_emotion_clip, - listener_3dmm_clip, - speaker_clip_length, - listener_clip_length, - ) = batch[:9] - listener_eeg_clip = listener_eeg_mask = None - if len(batch) > 9: - listener_eeg_clip, listener_eeg_mask = batch[9:11] - - (speaker_audio_clip, speaker_emotion_clip, speaker_3dmm_clip, - listener_video_clip, listener_emotion_clip, listener_3dmm_clip, - past_listener_emotion, past_listener_3dmm, motion_lengths, - listener_eeg_clip, listener_eeg_mask) = self.data_resample( - speaker_audio_clips=speaker_audio_clip, speaker_emotion_clips=speaker_emotion_clip, - speaker_3dmm_clips=speaker_3dmm_clip, listener_video_clips=listener_video_clip, - listener_emotion_clips=listener_emotion_clip, listener_3dmm_clips=listener_3dmm_clip, - speaker_seq_lengths=speaker_clip_length, listener_seq_lengths=listener_clip_length, - listener_eeg_clips=listener_eeg_clip, listener_eeg_masks=listener_eeg_mask) - - (speaker_audio_clip, # (78-d) - speaker_emotion_clip, # (25-d) - speaker_3dmm_clip, # (58-d) - listener_video_clip, - listener_emotion_clip, # (25-d) - ) = (speaker_audio_clip.to(device), - speaker_emotion_clip.to(device), - speaker_3dmm_clip.to(device), - listener_video_clip.to(device), - listener_emotion_clip.to(device)) - if listener_eeg_clip is not None: - listener_eeg_clip = listener_eeg_clip.to(device) - listener_eeg_mask = listener_eeg_mask.to(device) - batch_size = speaker_audio_clip.shape[0] - - # ---- scheduled sampling (online only): with prob ss_p, replace the GT past - # listener window with the model's OWN 1-step x̂₀ prediction of that window, so - # training matches the autoregressive inference distribution (closes exposure - # bias). Window-A speaker = [zero history | concurrent frames] (= the test's - # first-window condition); its GT target is exactly past_listener_emotion. ---- - past_for_B = past_listener_emotion - ss_p = 0.0 - if (self.task == "online" and past_listener_emotion is not None - and self._as_bool(self.trainer_cfg.get("scheduled_sampling", False))): - p_max = float(self.trainer_cfg.get("ss_p_max", 0.5)) - ramp = max(1, int(self.trainer_cfg.get("ss_ramp_epochs", self.epochs))) - ss_p = p_max * min(1.0, epoch / ramp) - if ss_p > 0.0: - bs_a = speaker_audio_clip.shape[0] - lw = int(self.trainer_cfg.window_size) # single window length (e.g. 30) - s_w_a = speaker_audio_clip.shape[1] - - def _winA(x): - hist = x.new_zeros(bs_a, s_w_a - lw, x.shape[-1]) - return torch.cat([hist, x[:, :lw]], dim=1) - - with torch.no_grad(): - out_A = model( - speaker_audio_input=_winA(speaker_audio_clip), - speaker_emotion_input=_winA(speaker_emotion_clip), - speaker_3dmm_input=_winA(speaker_3dmm_clip), - listener_emotion_input=past_listener_emotion[:, -lw:], # newest past window only - listener_eeg_input=None, listener_eeg_mask=None, - past_listener_emotion=None, - motion_length=None, - ) - dec_A = out_A.get("output_decoder", out_A) - x0_A = dec_A["prediction_emotion"].detach() # (bs, np, lw, 25) - npred = x0_A.shape[1] - gt_past_exp = past_listener_emotion.to(x0_A.device).repeat_interleave(npred, dim=0) # (bs*np, K*lw, 25) - # replace ONLY the newest window with the model's own x̂₀ (older windows stay GT) - self_full = gt_past_exp.clone() - self_full[:, -lw:, :] = x0_A.reshape(bs_a * npred, lw, x0_A.shape[-1]) - use_self = (torch.rand(bs_a, device=x0_A.device) < ss_p).repeat_interleave(npred) - past_for_B = torch.where(use_self[:, None, None], self_full, gt_past_exp) - - outputs = model( - speaker_audio_input=speaker_audio_clip, - speaker_emotion_input=speaker_emotion_clip, - speaker_3dmm_input=speaker_3dmm_clip, - listener_emotion_input=listener_emotion_clip, - listener_eeg_input=listener_eeg_clip, - listener_eeg_mask=listener_eeg_mask, - past_listener_emotion=past_for_B, - motion_length=motion_lengths, - ) - # outputs['prediction_emotion'].shape: [bs, k, l_w, 25] - # outputs['target_emotion'].shape: [bs, k, l_w, 25] - - output = criterion(outputs) - loss = output["loss_eeg"] if self.train_eeg_head_only else output["loss"] - if self.train_eeg_head_only and not loss.requires_grad: - raise RuntimeError( - "loss_eeg has no gradient. Check that EEG labels are enabled and prediction_eeg is returned." - ) - - iteration = batch_idx + len(data_loader) * epoch - if writer is not None: - writer.add_scalar("Train/loss", loss.data.item(), iteration) - writer.add_scalar("Train/loss_total", output["loss"].data.item(), iteration) - writer.add_scalar("Train/loss_prior", output["loss_prior"].data.item(), iteration) - writer.add_scalar("Train/loss_eeg", output["loss_eeg"].data.item(), iteration) - writer.add_scalar("Train/eeg_valid_ratio", output["eeg_valid_ratio"].data.item(), iteration) - # writer.add_scalar("Train/temporal_loss", temporal_loss.data.item(), iteration) - - whole_losses.update(loss.data.item(), batch_size) - prior_losses.update(output["loss_prior"].data.item(), batch_size) - au_rec_losses.update(output["loss_au"].data.item(), batch_size) - va_rec_losses.update(output["loss_va"].data.item(), batch_size) - em_rec_losses.update(output["loss_em"].data.item(), batch_size) - eeg_rec_losses.update(output["loss_eeg"].data.item(), batch_size) - eeg_valid_ratios.update(output["eeg_valid_ratio"].data.item(), batch_size) - - optimizer.zero_grad() - loss.backward() - optimizer.step() - if scheduler is not None and (epoch + 1) >= 5: - scheduler.step() - lr = get_lr(optimizer=optimizer) - if writer is not None: - writer.add_scalar("Train/lr", lr, epoch) - - return (whole_losses.avg, prior_losses.avg, au_rec_losses.avg, - va_rec_losses.avg, em_rec_losses.avg, eeg_rec_losses.avg, - eeg_valid_ratios.avg) - - def val_diffusion(self, model, val_loader, criterion, device): - whole_losses = AverageMeter() - prior_losses = AverageMeter() - au_rec_losses = AverageMeter() - va_rec_losses = AverageMeter() - em_rec_losses = AverageMeter() - eeg_rec_losses = AverageMeter() - eeg_valid_ratios = AverageMeter() - - model.eval() - for batch_idx, batch in enumerate(tqdm(val_loader)): - ( - speaker_audio_clip, - speaker_video_clip, - speaker_emotion_clip, - speaker_3dmm_clip, - listener_video_clip, - listener_emotion_clip, - listener_3dmm_clip, - speaker_clip_length, - listener_clip_length, - ) = batch[:9] - listener_eeg_clip = listener_eeg_mask = None - if len(batch) > 9: - listener_eeg_clip, listener_eeg_mask = batch[9:11] - - (speaker_audio_clip, speaker_emotion_clip, speaker_3dmm_clip, - listener_video_clip, listener_emotion_clip, listener_3dmm_clip, - past_listener_emotion, past_listener_3dmm, motion_lengths, - listener_eeg_clip, listener_eeg_mask) = self.data_resample( - speaker_audio_clips=speaker_audio_clip, speaker_emotion_clips=speaker_emotion_clip, - speaker_3dmm_clips=speaker_3dmm_clip, listener_video_clips=listener_video_clip, - listener_emotion_clips=listener_emotion_clip, listener_3dmm_clips=listener_3dmm_clip, - speaker_seq_lengths=speaker_clip_length, listener_seq_lengths=listener_clip_length, - listener_eeg_clips=listener_eeg_clip, listener_eeg_masks=listener_eeg_mask) - - (speaker_audio_clip, # (78-d) - speaker_emotion_clip, # (25-d) - speaker_3dmm_clip, # (58-d) - listener_video_clip, - listener_emotion_clip, # (25-d) - ) = (speaker_audio_clip.to(device), - speaker_emotion_clip.to(device), - speaker_3dmm_clip.to(device), - listener_video_clip.to(device), - listener_emotion_clip.to(device)) - if listener_eeg_clip is not None: - listener_eeg_clip = listener_eeg_clip.to(device) - listener_eeg_mask = listener_eeg_mask.to(device) - batch_size = speaker_audio_clip.shape[0] - - with torch.no_grad(): - outputs = model( - speaker_audio_input=speaker_audio_clip, - speaker_emotion_input=speaker_emotion_clip, - speaker_3dmm_input=speaker_3dmm_clip, - listener_emotion_input=listener_emotion_clip, - listener_eeg_input=listener_eeg_clip, - listener_eeg_mask=listener_eeg_mask, - past_listener_emotion=past_listener_emotion, - motion_length=motion_lengths, - ) - - output = criterion(outputs) - loss = output["loss_eeg"] if self.train_eeg_head_only else output["loss"] - whole_losses.update(loss.data.item(), batch_size) - prior_losses.update(output["loss_prior"].data.item(), batch_size) - au_rec_losses.update(output["loss_au"].data.item(), batch_size) - va_rec_losses.update(output["loss_va"].data.item(), batch_size) - em_rec_losses.update(output["loss_em"].data.item(), batch_size) - eeg_rec_losses.update(output["loss_eeg"].data.item(), batch_size) - eeg_valid_ratios.update(output["eeg_valid_ratio"].data.item(), batch_size) - - return (whole_losses.avg, prior_losses.avg, au_rec_losses.avg, - va_rec_losses.avg, em_rec_losses.avg, eeg_rec_losses.avg, - eeg_valid_ratios.avg) - - @staticmethod - def _eeg_targets_from_motion_lengths(listener_eeg, listener_eeg_mask, motion_lengths): - if listener_eeg is None or listener_eeg.numel() == 0: - return None, None - if listener_eeg_mask is None or listener_eeg_mask.numel() == 0: - listener_eeg_mask = torch.ones_like(listener_eeg) - - indices = [] - offset = 0 - total_length = listener_eeg.shape[0] - for motion_length in motion_lengths: - length = int(motion_length.item() if torch.is_tensor(motion_length) else motion_length) - last_idx = min(max(offset + max(length, 1) - 1, 0), total_length - 1) - indices.append(last_idx) - offset += max(length, 0) - if not indices: - return None, None - index_tensor = torch.tensor(indices, dtype=torch.long) - return ( - listener_eeg[index_tensor].unsqueeze(0).float(), - listener_eeg_mask[index_tensor].unsqueeze(0).float(), - ) - - def _frrea_render_sample(self, renderer, latent_embedder, pred_listener_emotion, - listener_video_clips, sample_idx, fake_dir, real_dir, - stride, shard_idx, batch_idx): - """Render one sample's generated (fake) + real listener frames for FRRea (FID). - - Renders only prediction #0. The full listener video is kept on CPU (test clips can - be very long); only the reference frame and the subsampled real frames are moved to - the GPU. Frame filenames are shard-unique so multiple GPU shards can write to the - same fake/real directories without collisions. - """ - import cv2 - lv = listener_video_clips[sample_idx] - if isinstance(lv, (list, tuple)): - lv = lv[0] - if not torch.is_tensor(lv) or lv.numel() == 0: - return - reference = lv[0].to(self.device) # (3, H, W); whole video stays on CPU - - emotion = pred_listener_emotion[0].to(self.device).float() # (clip_len, 25) - with torch.no_grad(): - coeff_3dmm = latent_embedder.decode_coeff(emotion) # (clip_len, 58) - - fake_np, real_np = renderer.render_frames_for_fid( - coeff_3dmm, reference, lv, fake_stride=stride) - - prefix = f"sh{shard_idx}_b{batch_idx}_s{sample_idx}" - for i in range(fake_np.shape[0]): - cv2.imwrite(os.path.join(fake_dir, f"{prefix}_f{i}.png"), fake_np[i]) - for i in range(real_np.shape[0]): - cv2.imwrite(os.path.join(real_dir, f"{prefix}_f{i}.png"), real_np[i]) - - def test(self): - stage = "test" - data_clamp = self.kwargs.pop("data_clamp") - eval_eeg = self._as_bool(self.trainer_cfg.get("eval_eeg", False)) - logger.info("Loading test data module") - test_loader = self.data_module.get_dataloader(stage=stage) - logger.info("Test data module loaded") - clip_len = self.trainer_cfg.clip_length - w = self.trainer_cfg.window_size - s_ratio = self.trainer_cfg.s_ratio - s_w = s_ratio * w - - model = instantiate(self.model_cfg.diff_model, - stage=stage, - latent_embedder=self.model_cfg.latent_embedder \ - if hasattr(self.model_cfg, "latent_embedder") else None, - audio_encoder=self.model_cfg.audio_encoder \ - if hasattr(self.model_cfg, "audio_encoder") else None, - **self.kwargs, - _recursive_=False) - model.to(self.device) - model.eval() - if eval_eeg: - if getattr(model, "eeg_head", None) is None: - raise RuntimeError( - "trainer.generic.eval_eeg=True but configs//model/motion_diffusion.yaml has no enabled eeg_head." - ) - eeg_ckpt_path = model.get_ckpt_path( - model.eeg_head, - runid="resume_runid", - epoch=None, - best=True, - last=False, - create_dir=False, - ) - if not os.path.exists(eeg_ckpt_path): - raise FileNotFoundError( - "trainer.generic.eval_eeg=True requires a trained EEGPredictionHead checkpoint. " - f"Missing: {eeg_ckpt_path}" - ) - - logger.info("Loading post processor") - post_processor = Processor(config_name=self.kwargs.pop("post_config_name"), - clip_len_test=self.kwargs.pop("post_clip_length"), - device=self.device,) - logger.info("Post processor loaded") - - GT_listener_emotions_all = [] - pred_listener_emotions_all = [] - input_speaker_emotions_all = [] - GT_listener_eeg_all = [] - pred_listener_eeg_all = [] - listener_eeg_mask_all = [] - - # ---- FRRea (FID) frame rendering setup (optional, gated by compute_frrea) ---- - compute_frrea = self._as_bool(self.kwargs.get("compute_frrea", False)) - frrea_renderer = frrea_latent_embedder = None - frrea_fake_dir = frrea_real_dir = None - frrea_stride = int(self.kwargs.get("frrea_stride", 30)) - frrea_shard_idx = int(os.environ.get("EVAL_SHARD_IDX", "0")) - if compute_frrea: - if self.renderer is None: - raise RuntimeError("compute_frrea=True requires trainer.renderer config.") - logger.info("Instantiating renderer for FRRea") - frrea_renderer = instantiate(self.renderer, device=self.device) - frrea_latent_embedder = model.diffusion_decoder.latent_embedder - tag = os.environ.get("FRREA_TAG") or str(self.kwargs.get("resume_runid", "run")) - base = os.path.join(to_absolute_path("frrea_frames"), str(self.task), tag) - frrea_fake_dir = os.path.join(base, "fake") - frrea_real_dir = os.path.join(base, "real") - os.makedirs(frrea_fake_dir, exist_ok=True) - os.makedirs(frrea_real_dir, exist_ok=True) - logger.info(f"FRRea frames -> {base} (shard {frrea_shard_idx}, stride {frrea_stride})") - - for batch_idx, batch in enumerate(tqdm(test_loader)): - ( - speaker_audio_clips, - speaker_video_clips, - speaker_emotion_clips, - speaker_3dmm_clips, - listener_video_clips, - listener_emotion_clips, - _, - speaker_seq_lengths, - listener_seq_lengths, - ) = batch[:9] - listener_eeg_clips = listener_eeg_masks = None - if len(batch) > 9: - listener_eeg_clips, listener_eeg_masks = batch[9:11] - if eval_eeg and listener_eeg_clips is None: - raise RuntimeError("trainer.generic.eval_eeg=True but the test dataloader did not return EEG labels.") - - # listener_emotion_clips: List: [[Tensor([l, d]), Tensor([l', d]), ...], ...] - for em in listener_emotion_clips: - GT_listener_emotions_all.append([em] if isinstance(em, torch.Tensor) else em) - input_speaker_emotions_all.extend(speaker_emotion_clips) - - clip_batch_size = 8 # in case too long data sequence - speaker_audios = [] - speaker_emotions = [] - speaker_3dmms = [] - motion_lengths = [] - sample_batch_size = [] - sample_eeg_targets = [] - sample_eeg_masks = [] - eeg_clips = listener_eeg_clips if eval_eeg else [None] * len(speaker_audio_clips) - eeg_masks = listener_eeg_masks if eval_eeg else [None] * len(speaker_audio_clips) - - for (speaker_audio_clip, speaker_emotion_clip, speaker_3dmm_clip, - speaker_seq_length, listener_eeg_clip, listener_eeg_mask) in zip( - speaker_audio_clips, speaker_emotion_clips, speaker_3dmm_clips, - speaker_seq_lengths, eeg_clips, eeg_masks): - length = int(speaker_seq_length.item() if torch.is_tensor(speaker_seq_length) else speaker_seq_length) - - # Align all speaker clips to exactly `length` frames - # (.npy feature files may have slightly different frame counts than the video) - def _align_clip(clip, tgt_len): - if clip.dim() < 1 or clip.shape[0] == tgt_len: - return clip - if clip.shape[0] > tgt_len: - return clip[:tgt_len] - return torch.cat([clip, torch.zeros(tgt_len - clip.shape[0], clip.shape[-1])], dim=0) - - speaker_audio_clip = _align_clip(speaker_audio_clip, length) - speaker_emotion_clip = _align_clip(speaker_emotion_clip, length) - speaker_3dmm_clip = _align_clip(speaker_3dmm_clip, length) - - if self.task == "offline": - remain_length = length % clip_len - b = max(math.ceil(length / clip_len), 1) - final_length = remain_length if remain_length != 0 else clip_len - lengths = torch.tensor([clip_len] * (b - 1) + [final_length]) - sample_batch_size.append(b) - pad_length = b * clip_len - length - - speaker_audio_clip = torch.cat((speaker_audio_clip, - torch.zeros( - size=(pad_length, speaker_audio_clip.shape[-1]))), - dim=0) - speaker_audio_clip = rearrange(speaker_audio_clip, '(b l) d -> b l d', b=b) - - speaker_emotion_clip = torch.cat((speaker_emotion_clip, - torch.zeros(size=(pad_length, - speaker_emotion_clip.shape[-1]))), dim=0) - speaker_emotion_clip = rearrange(speaker_emotion_clip, '(b l) d -> b l d', b=b) - - speaker_3dmm_clip = torch.cat((speaker_3dmm_clip, - torch.zeros( - size=(pad_length, speaker_3dmm_clip.shape[-1]))), - dim=0) - speaker_3dmm_clip = rearrange(speaker_3dmm_clip, '(b l) d -> b l d', b=b) - - speaker_audios.append(speaker_audio_clip) - speaker_emotions.append(speaker_emotion_clip) - speaker_3dmms.append(speaker_3dmm_clip) - motion_lengths.append(lengths) - if eval_eeg: - eeg_target, eeg_mask = self._eeg_targets_from_motion_lengths( - listener_eeg_clip, listener_eeg_mask, lengths) - sample_eeg_targets.append(eeg_target) - sample_eeg_masks.append(eeg_mask) - - else: # online task - num_windows = math.ceil(length / w) - sample_batch_size.append(num_windows) - - speaker_audio_clip = torch.cat( - (torch.zeros(size=((s_w - w), speaker_audio_clip.shape[-1])), - speaker_audio_clip, - torch.zeros(size=((num_windows * w - length), speaker_audio_clip.shape[-1]))), dim=0) - speaker_emotion_clip = torch.cat( - (torch.zeros(size=((s_w - w), speaker_emotion_clip.shape[-1])), - speaker_emotion_clip, - torch.zeros(size=((num_windows * w - length), speaker_emotion_clip.shape[-1]))), dim=0) - speaker_3dmm_clip = torch.cat( - (torch.zeros(size=((s_w - w), speaker_3dmm_clip.shape[-1])), - speaker_3dmm_clip, - torch.zeros(size=((num_windows * w - length), speaker_3dmm_clip.shape[-1]))), dim=0) - - motion_length_list = [] - speaker_audio_clip_list = [] - speaker_emotion_clip_list = [] - speaker_3dmm_clip_list = [] - for i in range(num_windows): - motion_length_list.append(w) if i < num_windows - 1 else motion_length_list.append( - length - i * w) - speaker_audio_clip_list.append(speaker_audio_clip[i*w: i*w + s_w]) - speaker_emotion_clip_list.append(speaker_emotion_clip[i*w: i*w + s_w]) - speaker_3dmm_clip_list.append(speaker_3dmm_clip[i*w: i*w + s_w]) - - motion_length = torch.tensor(motion_length_list) - speaker_audio_clip = torch.stack(speaker_audio_clip_list, dim=0) - speaker_emotion_clip = torch.stack(speaker_emotion_clip_list, dim=0) - speaker_3dmm_clip = torch.stack(speaker_3dmm_clip_list, dim=0) - - motion_lengths.append(motion_length) - speaker_audios.append(speaker_audio_clip) - speaker_emotions.append(speaker_emotion_clip) - speaker_3dmms.append(speaker_3dmm_clip) - if eval_eeg: - eeg_target, eeg_mask = self._eeg_targets_from_motion_lengths( - listener_eeg_clip, listener_eeg_mask, motion_length) - sample_eeg_targets.append(eeg_target) - sample_eeg_masks.append(eeg_mask) - - motion_lengths = torch.cat(motion_lengths, dim=0) - speaker_audios = torch.cat(speaker_audios, dim=0) - speaker_emotions = torch.cat(speaker_emotions, dim=0) - speaker_3dmms = torch.cat(speaker_3dmms, dim=0) - sample_batch_size = torch.tensor(sample_batch_size) - - pred_listener_emotions = [] - pred_listener_eegs = [] - frrea_skip = False - all_batch_size = speaker_audios.shape[0] - - if self.task == "online" and getattr(self, "online_autoregressive", True): - # Autoregressive cross-window continuity: the listener windows of one - # sample are generated sequentially, each conditioned on the PREVIOUS - # window's own prediction (per-prediction past_listener). This restores - # temporal continuity across the stitched sequence instead of generating - # every 30-frame window independently. Window 0 has no past (None). - win_bounds = torch.cat( - (torch.tensor([0]), torch.cumsum(sample_batch_size, dim=0)), dim=0) - K = int(self.trainer_cfg.get("n_past_win", 1)) - lw_t = int(self.trainer_cfg.window_size) - for si in range(len(sample_batch_size)): - a, b = int(win_bounds[si]), int(win_bounds[si + 1]) - buf = [] # rolling buffer of own generated windows, each (num_preds, lw, 25) - for wi in range(a, b): - spk_a = speaker_audios[wi:wi + 1].to(self.device) - spk_e = speaker_emotions[wi:wi + 1].to(self.device) - spk_3 = speaker_3dmms[wi:wi + 1].to(self.device) - ml = motion_lengths[wi:wi + 1].to(self.device) - # past = last K generated windows, front-padded with zeros to K*lw - # (matches training); None for the very first window. - if len(buf) == 0: - past = None - else: - cat_buf = torch.cat(buf[-K:], dim=1) # (np, k*lw, 25) - need = K * lw_t - cat_buf.shape[1] - if need > 0: - cat_buf = torch.cat( - [cat_buf.new_zeros(cat_buf.shape[0], need, cat_buf.shape[2]), cat_buf], dim=1) - past = cat_buf - try: - with torch.no_grad(): - outputs = model( - speaker_audio_input=spk_a, - speaker_emotion_input=spk_e, - speaker_3dmm_input=spk_3, - motion_length=ml, - past_listener_emotion=past, - ) - except RuntimeError as e: - if compute_frrea: - logger.warning(f"FRRea: skipping sample (batch {batch_idx}) due to: {e}") - torch.cuda.empty_cache() - frrea_skip = True - break - raise - pred = outputs["prediction_emotion"] # (1, num_preds, l_w, 25) - pred_listener_emotions.append(pred.detach().cpu()) - if eval_eeg: - if "prediction_eeg" not in outputs: - raise RuntimeError("trainer.generic.eval_eeg=True but the model did not return prediction_eeg.") - pred_listener_eegs.append(outputs["prediction_eeg"].detach().cpu()) - buf.append(pred[0].detach()) # (num_preds, lw, 25) - if len(buf) > K: - buf.pop(0) - if frrea_skip: - break - else: - for i in range(math.ceil(all_batch_size / clip_batch_size)): - speaker_audio_clip = speaker_audios[i * clip_batch_size: (i + 1) * clip_batch_size] - speaker_emotion_clip = speaker_emotions[i * clip_batch_size: (i + 1) * clip_batch_size] - speaker_3dmm_clip = speaker_3dmms[i * clip_batch_size: (i + 1) * clip_batch_size] - motion_length = motion_lengths[i * clip_batch_size: (i + 1) * clip_batch_size] - - (speaker_audio_clip, - speaker_emotion_clip, - speaker_3dmm_clip) = ( - speaker_audio_clip.to(self.device), - speaker_emotion_clip.to(self.device), - speaker_3dmm_clip.to(self.device)) - # speaker_audio_clip: (bsz, s_w, d_audio) - # speaker_emotion_clip: (bsz, s_w, d_emotion) - # speaker_3dmm_clip: (bsz, s_w, d_3dmm) - - try: - with torch.no_grad(): - outputs = model( - speaker_audio_input=speaker_audio_clip, - speaker_emotion_input=speaker_emotion_clip, - speaker_3dmm_input=speaker_3dmm_clip, - motion_length=motion_length, - ) - except RuntimeError as e: - # A single pathological sample (e.g. a very long clip causing GPU OOM) - # should not abort the whole FRRea render run; skip it. - if compute_frrea: - logger.warning(f"FRRea: skipping sample (batch {batch_idx}) due to: {e}") - torch.cuda.empty_cache() - frrea_skip = True - break - raise - - pred_listener_emotions.append(outputs["prediction_emotion"].detach().cpu()) - if eval_eeg: - if "prediction_eeg" not in outputs: - raise RuntimeError("trainer.generic.eval_eeg=True but the model did not return prediction_eeg.") - pred_listener_eegs.append(outputs["prediction_eeg"].detach().cpu()) - if frrea_skip: - torch.cuda.empty_cache() - continue - pred_listener_emotions = torch.cat(pred_listener_emotions, dim=0) # (L', num_preds, l_w, 25) - pred_listener_eegs = torch.cat(pred_listener_eegs, dim=0) if eval_eeg else None - - bounds = torch.cat((torch.tensor([0]), torch.cumsum(sample_batch_size, dim=0)), dim=0) - intervals = list(zip(bounds[:-1], bounds[1:])) - for sample_idx, (l, r) in enumerate(intervals): - pred_listener_emotion = pred_listener_emotions[l:r] # (b', num_preds, l_w, 25) - motion_length = motion_lengths[l:r] - clip_length = int(torch.sum(motion_length, dim=0, keepdim=False).item()) - pred_listener_emotion = rearrange(pred_listener_emotion, - 'b n w d -> n (b w) d')[:, :clip_length] - - if data_clamp: - pred_listener_emotion[:, :, :15] = torch.round(pred_listener_emotion[:, :, :15]) - - pred_listener_emotions_all.append(pred_listener_emotion) - if compute_frrea: - self._frrea_render_sample( - frrea_renderer, frrea_latent_embedder, pred_listener_emotion, - listener_video_clips, sample_idx, frrea_fake_dir, frrea_real_dir, - frrea_stride, frrea_shard_idx, batch_idx) - if eval_eeg: - pred_listener_eeg = rearrange(pred_listener_eegs[l:r], 'b n d -> n b d') - pred_listener_eeg_all.append(pred_listener_eeg) - GT_listener_eeg_all.append(sample_eeg_targets[sample_idx]) - listener_eeg_mask_all.append(sample_eeg_masks[sample_idx]) - - # pred_listener_emotions_all - # List: 750 [Tensor([num_preds, l, 25]), Tensor([num_preds, l', 25]), ...] - # GT_listener_emotions_all - # List: 750 [List: [(l', 25), (l'', 25), ...], List: [(l''', 25), (l'''', 25)], ...] - if len(pred_listener_emotions_all): - GT_listener_emotions_all = post_processor.forward( - prediction_list=pred_listener_emotions_all, - target_list=GT_listener_emotions_all,) - # GT_listener_emotions_all - # List: 750 [Tensor([num_preds, l, 25]), Tensor([num_preds, l', 25]), ...] - - try: - result_dict = {'GT': GT_listener_emotions_all, 'PRED': pred_listener_emotions_all} - if eval_eeg: - result_dict.update({ - 'GT_EEG': GT_listener_eeg_all, - 'PRED_EEG': pred_listener_eeg_all, - 'EEG_MASK': listener_eeg_mask_all, - }) - torch.save(result_dict, f'results.pt') - print("Successfully saved Tensor List") - except Exception: - print("Failed to save Tensor List") - - if compute_frrea: - # FRRea render run uses num_preds=1, for which the per-sample emotion metrics - # (e.g. S-MSE diversity) are undefined; skip them and only emit the frames. - results = {} - else: - results = compute_metrics( - input_speaker_emotions_all, - pred_listener_emotions_all, - GT_listener_emotions_all, - ) - if eval_eeg: - results.update(compute_eeg_metrics( - pred_listener_eeg_all, - GT_listener_eeg_all, - listener_eeg_mask_all, - )) - logger.info(results) - if compute_frrea: - logger.info( - "FRRea frames written. After all shards finish, compute FID with:\n" - f" python -m framework.metrics.FID --fake {frrea_fake_dir} --real {frrea_real_dir}" - ) diff --git a/personalised/code/trainer/motion_diffusion_dp.py b/personalised/code/trainer/motion_diffusion_dp.py deleted file mode 100644 index 501fe68455c227059a643ec97121051e6e6a6f8c..0000000000000000000000000000000000000000 --- a/personalised/code/trainer/motion_diffusion_dp.py +++ /dev/null @@ -1,119 +0,0 @@ -""" -DataParallel wrapper for motion_diffusion Trainer. -Overrides main_diffusion() to use nn.DataParallel when multiple GPUs are available. -All other logic (test, val, data_resample, etc.) is inherited unchanged. -""" -import torch -import torch.nn as nn -import logging -from hydra.utils import instantiate -from framework.utils.util import from_pretrained_checkpoint - -from trainer.motion_diffusion import Trainer as _BaseTrainer - -logger = logging.getLogger(__name__) - - -class Trainer(_BaseTrainer): - - def main_diffusion(self, stage): - if self.train_eeg_head_only and self.resumed_training: - raise ValueError( - "train_eeg_head_only=True should be launched with resume=false, " - "so the optimizer and old EEG head checkpoint are not restored." - ) - - model = instantiate(self.model_cfg.diff_model, - stage=stage, - resumed_training=self.resumed_training, - latent_embedder=self.model_cfg.latent_embedder \ - if hasattr(self.model_cfg, "latent_embedder") else None, - audio_encoder=self.model_cfg.audio_encoder \ - if hasattr(self.model_cfg, "audio_encoder") else None, - **self.kwargs, - _recursive_=False) - model.to(self.device) - self._load_pretrained_motion_diffusion(model) - - # --- DataParallel --- - n_gpus = torch.cuda.device_count() - if n_gpus > 1: - logger.info(f"Using DataParallel across {n_gpus} GPUs") - model = nn.DataParallel(model) - model_raw = model.module if isinstance(model, nn.DataParallel) else model - # -------------------- - - optimizer_params = model_raw.parameters() - if self.train_eeg_head_only: - model_raw.freeze_except_eeg_head() - trainable_params = [p for p in model_raw.parameters() if p.requires_grad] - trainable_names = [n for n, p in model_raw.named_parameters() if p.requires_grad] - frozen_count = self._count_parameters( - p for p in model_raw.parameters() if not p.requires_grad - ) - trainable_count = self._count_parameters(trainable_params) - if len(trainable_params) == 0: - raise RuntimeError("No trainable parameters found for EEG head-only training.") - optimizer_params = trainable_params - logger.info( - "EEG head-only training enabled. " - f"Trainable parameters: {trainable_count}; frozen parameters: {frozen_count}" - ) - logger.info(f"Trainable parameter tensors: {trainable_names}") - - optimizer = instantiate(self.optim_cfg, lr=self.trainer_cfg.lr, params=optimizer_params) - if self.resumed_training: - checkpoint_path = model_raw.get_ckpt_path( - model_raw.diffusion_decoder.model, runid="resume_runid", last=True - ) - best_validation_loss, self.start_epoch = ( - from_pretrained_checkpoint(checkpoint_path, optimizer, self.device) - ) - logger.info(f"Resume training from epoch {self.start_epoch}") - else: - best_validation_loss = float('inf') - print(f"Best validation loss: {best_validation_loss}") - - scheduler = instantiate(self.kwargs.pop("scheduler"), optimizer, len(self.train_loader)) - selected_loss_name = "loss_eeg" if self.train_eeg_head_only else "diff_loss" - - for epoch in range(self.start_epoch, self.epochs): - if self.train_eeg_head_only: - model_raw.set_eeg_head_train_mode() - else: - model.train() - - diffusion_loss, prior_loss, au_rec_loss, va_rec_loss, em_rec_loss, eeg_rec_loss, eeg_valid_ratio = ( - self.train_diffusion(model, self.train_loader, optimizer, scheduler, - self.criterion, epoch, self.writer, self.device)) - logging.info(f"Epoch: {epoch + 1} train_{selected_loss_name}: {diffusion_loss:.5f} " - f"prior_loss: {prior_loss:.5f} au_rec_loss: {au_rec_loss:.5f}" - f" va_rec_loss: {va_rec_loss:.5f} em_rec_loss: {em_rec_loss:.5f}" - f" eeg_rec_loss: {eeg_rec_loss:.5f} eeg_valid_ratio: {eeg_valid_ratio:.5f}") - if self.writer is not None: - self.writer.add_scalar("Epoch/train_loss", diffusion_loss, epoch + 1) - - if (epoch + 1) % self.val_period == 0: - diffusion_loss, prior_loss, au_rec_loss, va_rec_loss, em_rec_loss, eeg_rec_loss, eeg_valid_ratio = ( - self.val_diffusion(model, self.val_loader, self.criterion, self.device)) - logging.info(f"Epoch: {epoch + 1} val_{selected_loss_name}: {diffusion_loss:.5f} " - f"prior_loss: {prior_loss:.5f} au_rec_loss: {au_rec_loss:.5f}" - f" va_rec_loss: {va_rec_loss:.5f} em_rec_loss: {em_rec_loss:.5f}" - f" eeg_rec_loss: {eeg_rec_loss:.5f} eeg_valid_ratio: {eeg_valid_ratio:.5f}") - if self.writer is not None: - self.writer.add_scalar("Epoch/val_loss", diffusion_loss, epoch + 1) - self.writer.add_scalar("Epoch/val_prior_loss", prior_loss, epoch + 1) - self.writer.add_scalar("Epoch/val_au_rec_loss", au_rec_loss, epoch + 1) - self.writer.add_scalar("Epoch/val_va_rec_loss", va_rec_loss, epoch + 1) - self.writer.add_scalar("Epoch/val_em_rec_loss", em_rec_loss, epoch + 1) - - if diffusion_loss < best_validation_loss: - best_validation_loss = diffusion_loss - logging.info( - f"New best {selected_loss_name} ({best_validation_loss:.5f}) at epoch {epoch + 1}, " - f"saving checkpoint" - ) - model_raw.save_ckpt(optimizer, best=True, epoch=(epoch + 1), best_loss=best_validation_loss) - - model_raw.save_ckpt(optimizer, epoch=(epoch + 1), best_loss=best_validation_loss) - model_raw.save_ckpt(optimizer, last=True, epoch=(epoch + 1), best_loss=best_validation_loss) diff --git a/personalised/code/trainer/motion_transvae.py b/personalised/code/trainer/motion_transvae.py deleted file mode 100644 index df9b20fd8ed9d8e4ebcad41726cde555c6bc316a..0000000000000000000000000000000000000000 --- a/personalised/code/trainer/motion_transvae.py +++ /dev/null @@ -1,788 +0,0 @@ -import math -from pathlib import Path -from typing import List - -import hydra -from hydra.utils import instantiate -from omegaconf import DictConfig -import os -import torch -import torch.nn as nn -import torch.optim as optim -import argparse -from torch.utils.tensorboard import SummaryWriter -from tqdm import tqdm -import logging -from framework.modules.post_processor import Processor -from framework.utils.compute_metrics import compute_eeg_metrics, compute_metrics -from framework.utils.losses import div_loss -from framework.utils.util import AverageMeter, from_pretrained_checkpoint - -os.environ["NUMEXPR_MAX_THREADS"] = '16' -logger = logging.getLogger(__name__) - - -class Trainer: - def __init__(self, - resumed_training: bool = False, - renderer: DictConfig = None, - model: DictConfig = None, - criterion: DictConfig = None, - **kwargs): - - self.renderer_cfg = renderer - self.model_cfg = model - self.criterion_cfg = criterion - self.resumed_training = resumed_training - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - self.lr = kwargs.pop('lr') - self.optim_cfg = kwargs.pop("optim") - self.epochs = kwargs.pop("epochs") - self.gpu_ids = kwargs.pop("gpu_ids") - self.j = kwargs.pop("j") - self.max_seq_len = kwargs.pop("max_seq_len") - self.window_size = kwargs.pop("window_size") - self.div_p = kwargs.pop("div_p") - self.task = kwargs.pop("task") - self.train_eeg_head_only = self._as_bool(kwargs.pop("train_eeg_head_only", False)) - self.eval_eeg = self._as_bool(kwargs.pop("eval_eeg", False)) - self.pretrained_model_checkpoint = kwargs.pop("pretrained_model_checkpoint", "") - self.num_preds = kwargs.pop("num_preds", 10) - self.save_results = self._as_bool(kwargs.pop("save_results", True)) - self.eval_facial_metrics = self._as_bool(kwargs.pop("eval_facial_metrics", True)) - self.eval_eeg_metrics = self._as_bool(kwargs.pop("eval_eeg_metrics", True)) - self.metric_threads = int(kwargs.pop("metric_threads", 1)) - self.eval_clip_batch_size = max(int(kwargs.pop("eval_clip_batch_size", 1)), 1) - self.kwargs = kwargs - - @staticmethod - def _as_bool(value): - if isinstance(value, str): - return value.lower() in {"1", "true", "yes", "y"} - return bool(value) - - @staticmethod - def _count_parameters(parameters): - return sum(parameter.numel() for parameter in parameters) - - @staticmethod - def _resolve_checkpoint_path(path): - if path is None or str(path).strip() == "": - return None - path = str(path) - if os.path.isabs(path): - return path - return hydra.utils.to_absolute_path(path) - - def get_ckpt_path(self, model, runid="current_runid", epoch=None, best=False, last=False, create_dir=True): - ckpt_dir = Path(hydra.utils.to_absolute_path(self.kwargs.get("ckpt_dir"))) - run_id = Path(self.kwargs.get(runid)) - ckpt_dir = str(ckpt_dir / run_id / model.get_model_name()) - if create_dir: - os.makedirs(ckpt_dir, exist_ok=True) - ckpt_path = None - if epoch is not None: - ckpt_path = os.path.join(ckpt_dir, f"checkpoint_{epoch}.pth") - if best: - ckpt_path = os.path.join(ckpt_dir, "checkpoint_best.pth") - if last: - ckpt_path = os.path.join(ckpt_dir, "checkpoint_last.pth") - assert ckpt_path is not None, "No checkpoint path is provided." - return ckpt_path - - @staticmethod - def _checkpoint_state_dict(checkpoint): - return checkpoint["state_dict"] if isinstance(checkpoint, dict) and "state_dict" in checkpoint else checkpoint - - def _load_model_checkpoint(self, model, checkpoint_path, strict=True): - checkpoint = torch.load(checkpoint_path, map_location="cpu") - state_dict = self._checkpoint_state_dict(checkpoint) - result = model.load_state_dict(state_dict, strict=strict) - model.to(self.device) - missing = getattr(result, "missing_keys", []) - unexpected = getattr(result, "unexpected_keys", []) - if missing: - logger.warning(f"Missing keys while loading {checkpoint_path}: {missing}") - if unexpected: - logger.warning(f"Unexpected keys while loading {checkpoint_path}: {unexpected}") - print(f"Successfully load model checkpoint: {checkpoint_path}") - if isinstance(checkpoint, dict): - return checkpoint.get("best_loss", float("inf")), checkpoint.get("epoch", 0) - return float("inf"), 0 - - def _load_pretrained_transvae(self, model): - checkpoint_path = self._resolve_checkpoint_path(self.pretrained_model_checkpoint) - if not self.train_eeg_head_only and checkpoint_path is None: - return - if self.train_eeg_head_only and checkpoint_path is None: - raise ValueError( - "train_eeg_head_only=True requires trainer.pretrained_model_checkpoint. " - "Point it to a pretrained TransformerVAE checkpoint and launch with resume=false." - ) - if not os.path.exists(checkpoint_path): - raise FileNotFoundError(f"Missing pretrained TransformerVAE checkpoint: {checkpoint_path}") - self._load_model_checkpoint(model, checkpoint_path, strict=False) - logger.info(f"Loaded pretrained TransformerVAE checkpoint: {checkpoint_path}") - - def _save_checkpoints(self, model, optimizer, epoch=None, best=False, last=False, best_loss=float("inf")): - checkpoint = { - "epoch": epoch, - "best_loss": best_loss, - "state_dict": model.state_dict(), - "optimizer": optimizer.state_dict(), - } - torch.save(checkpoint, self.get_ckpt_path(model, epoch=epoch, best=best, last=last)) - - if getattr(model, "eeg_head", None) is not None: - eeg_checkpoint = { - "epoch": epoch, - "best_loss": best_loss, - "state_dict": model.eeg_head.state_dict(), - "optimizer": optimizer.state_dict(), - } - torch.save(eeg_checkpoint, self.get_ckpt_path(model.eeg_head, epoch=epoch, best=best, last=last)) - - def set_data_module(self, data_module): - self.data_module = data_module - - def data_resample(self, speaker_audio_clips, speaker_video_clips, speaker_emotion_clips, speaker_3dmm_clips, - listener_emotion_clips, listener_3dmm_clips, speaker_seq_lengths, listener_seq_lengths, - listener_eeg_clips=None, listener_eeg_masks=None): - speaker_audios = [audio[:L] for audio, L in zip(speaker_audio_clips, speaker_seq_lengths)] - speaker_videos = [video[:L] for video, L in zip(speaker_video_clips, speaker_seq_lengths)] - speaker_emotions = [emo[:L] for emo, L in zip(speaker_emotion_clips, speaker_seq_lengths)] - speaker_3dmm = [param[:L] for param, L in zip(speaker_3dmm_clips, speaker_seq_lengths)] - listener_emotions = [emo[:L] for emo, L in zip(listener_emotion_clips, listener_seq_lengths)] - listener_3dmm = [param[:L] for param, L in zip(listener_3dmm_clips, listener_seq_lengths)] - listener_eegs = listener_eeg_masks_out = None - if listener_eeg_clips is not None and listener_eeg_masks is not None: - listener_eegs = [eeg[:L] for eeg, L in zip(listener_eeg_clips, listener_seq_lengths)] - listener_eeg_masks_out = [mask[:L] for mask, L in zip(listener_eeg_masks, listener_seq_lengths)] - return ( - speaker_audios, - speaker_videos, - speaker_emotions, - speaker_3dmm, - listener_emotions, - listener_3dmm, - listener_eegs, - listener_eeg_masks_out, - ) - - def fit(self): - """ - # relative directory - root_dir = save/${trainer.task_name}/${data.data_name}/${folder_name} - # absolute directory - saving_dir = Path(hydra.utils.to_absolute_path(root_dir)) - # get saving path - saving_path = str(saving_dir / ...) - """ - stage = "fit" - - logger.info("Loading data module") - self.train_loader, self.val_loader = ( - self.data_module.get_dataloader(stage=stage)) - logger.info("Data module loaded") - - logger.info("Loading criterion") - self.criterion = instantiate(self.criterion_cfg) - logger.info("Criterion loaded") - - self.main() - - def main(self): - if self.train_eeg_head_only and self.resumed_training: - raise ValueError( - "train_eeg_head_only=True should be launched with resume=false, " - "so the optimizer and old EEG head checkpoint are not restored." - ) - - model = instantiate(self.model_cfg, _recursive_=False) - model.to(self.device) - self._load_pretrained_transvae(model) - - optimizer_params = model.parameters() - if self.train_eeg_head_only: - model.freeze_except_eeg_head() - trainable_params = [parameter for parameter in model.parameters() if parameter.requires_grad] - trainable_names = [name for name, parameter in model.named_parameters() if parameter.requires_grad] - frozen_count = self._count_parameters( - parameter for parameter in model.parameters() if not parameter.requires_grad - ) - trainable_count = self._count_parameters(trainable_params) - if len(trainable_params) == 0: - raise RuntimeError("No trainable parameters found for EEG head-only training.") - optimizer_params = trainable_params - logger.info( - "EEG head-only training enabled. " - f"Trainable parameters: {trainable_count}; frozen parameters: {frozen_count}" - ) - logger.info(f"Trainable parameter tensors: {trainable_names}") - print( - "EEG head-only training enabled. " - f"Trainable parameters: {trainable_count}; frozen parameters: {frozen_count}" - ) - print(f"Trainable parameter tensors: {trainable_names}") - - optimizer = instantiate(self.optim_cfg, lr=self.lr, params=optimizer_params) - - tb_dir = hydra.utils.to_absolute_path( - os.path.join("tb_logs", self.kwargs.get("current_runid", "default"))) - writer = SummaryWriter(log_dir=tb_dir) - logger.info(f"TensorBoard log dir: {tb_dir}") - - if self.resumed_training: - checkpoint_path = self.get_ckpt_path(model, runid="resume_runid", last=True) - from_pretrained_checkpoint(checkpoint_path, optimizer, self.device) - lowest_val_loss, start_epoch = self._load_model_checkpoint(model, checkpoint_path, strict=False) - logger.info(f"Resume training from epoch {start_epoch}") - else: - start_epoch = 0 - lowest_val_loss = float('inf') - print(f"Best validation loss: {lowest_val_loss}") - - for epoch in range(start_epoch, self.epochs): - train_loss, rec_loss, rec_emo_loss, rec_param_loss, kld_loss, div_loss, eeg_loss, eeg_valid_ratio = ( - self.train(model, optimizer) - ) - logger.info("Epoch: {} train_loss: {:.5f} rec_all_loss: {:.5f} rec_emo_loss: {:.5f} " - "rec_parma_loss: {:.5f} kld_loss: {:.5f} div_loss: {:.5f} " - "eeg_loss: {:.5f} eeg_valid_ratio: {:.5f}" - .format(epoch + 1, train_loss, rec_loss, rec_emo_loss, rec_param_loss, kld_loss, div_loss, - eeg_loss, eeg_valid_ratio)) - writer.add_scalar("Train/loss", train_loss, epoch + 1) - writer.add_scalar("Train/rec_loss", rec_loss, epoch + 1) - writer.add_scalar("Train/rec_emo_loss", rec_emo_loss, epoch + 1) - writer.add_scalar("Train/rec_param_loss", rec_param_loss, epoch + 1) - writer.add_scalar("Train/kld_loss", kld_loss, epoch + 1) - writer.add_scalar("Train/div_loss", div_loss, epoch + 1) - - if (epoch + 1) % 5 == 0: - val_loss, rec_loss, rec_emo_loss, rec_param_loss, kld_loss, eeg_loss, eeg_valid_ratio = self.val(model) - logger.info("Epoch: {} val_loss: {:.5f} rec_all_loss: {:.5f} rec_emo_loss: {:.5f} " - "rec_param_loss: {:.5f} kld_loss: {:.5f} " - "eeg_loss: {:.5f} eeg_valid_ratio: {:.5f}" - .format(epoch + 1, val_loss, rec_loss, rec_emo_loss, rec_param_loss, kld_loss, - eeg_loss, eeg_valid_ratio)) - writer.add_scalar("Val/loss", val_loss, epoch + 1) - writer.add_scalar("Val/rec_emo_loss", rec_emo_loss, epoch + 1) - writer.add_scalar("Val/rec_param_loss", rec_param_loss, epoch + 1) - writer.add_scalar("Val/kld_loss", kld_loss, epoch + 1) - - if val_loss < lowest_val_loss: - lowest_val_loss = val_loss - ckpt_path = self.get_ckpt_path(model, best=True) - logger.info(f"Saving best checkpoint, val_loss: {lowest_val_loss:.5f}, ckpt_path: {ckpt_path}") - self._save_checkpoints( - model, optimizer, best=True, epoch=(epoch + 1), best_loss=lowest_val_loss) - - self._save_checkpoints(model, optimizer, epoch=(epoch + 1), best_loss=lowest_val_loss) - self._save_checkpoints(model, optimizer, last=True, epoch=(epoch + 1), best_loss=lowest_val_loss) - - writer.close() - - # Train - def train(self, model, optimizer): - losses = AverageMeter() - rec_losses = AverageMeter() - rec_emo_losses = AverageMeter() - rec_param_losses = AverageMeter() - kld_losses = AverageMeter() - div_losses = AverageMeter() - eeg_losses = AverageMeter() - eeg_valid_ratios = AverageMeter() - - if self.train_eeg_head_only: - model.set_eeg_head_train_mode() - else: - model.train() - for batch_idx, batch in enumerate(tqdm(self.train_loader)): - ( - speaker_audio_clip, - speaker_video_clip, - speaker_emotion_clip, - speaker_3dmm_clip, - _, - listener_emotion, - listener_3dmm, - speaker_clip_length, - listener_clip_length, - ) = batch[:9] - listener_eeg_clip = listener_eeg_mask = None - if len(batch) > 9: - listener_eeg_clip, listener_eeg_mask = batch[9:11] - - if self.model_cfg.task == 'offline': - (speaker_audio_clip, speaker_video_clip, speaker_emotion_clip, speaker_3dmm_clip, - listener_emotion, listener_3dmm, listener_eeg_clip, listener_eeg_mask) = self.data_resample( - speaker_audio_clip, - speaker_video_clip, - speaker_emotion_clip, - speaker_3dmm_clip, - listener_emotion, - listener_3dmm, - speaker_clip_length, - listener_clip_length, - listener_eeg_clips=listener_eeg_clip, - listener_eeg_masks=listener_eeg_mask, - ) - optimizer.zero_grad() - listener_3dmm_out, listener_emotion_out, distribution, eeg_outputs = model( - speaker_video_clip, - speaker_audio_clip, - motion_lengths=speaker_clip_length, - speaker_emotion=speaker_emotion_clip, - speaker_3dmm=speaker_3dmm_clip, - listener_eeg_input=listener_eeg_clip, - listener_eeg_mask=listener_eeg_mask, - return_eeg_outputs=True, - ) - - loss, rec_loss, rec_emo_loss, rec_param_loss, kld_loss, eeg_loss, eeg_valid_ratio = self.criterion( - listener_emotion, - listener_3dmm, - listener_emotion_out, - listener_3dmm_out, - distribution, - prediction_eeg=eeg_outputs.get("prediction_eeg"), - target_eeg=eeg_outputs.get("target_eeg"), - target_eeg_mask=eeg_outputs.get("target_eeg_mask"), - ) - if self.train_eeg_head_only: - loss = eeg_loss - d_loss = loss.new_tensor(0.0) - if not loss.requires_grad: - raise RuntimeError( - "loss_eeg has no gradient. Check that EEG labels are enabled and prediction_eeg is returned." - ) - else: - with torch.no_grad(): - listener_3dmm_out_, listener_emotion_out_, _ = model( - speaker_video_clip, - speaker_audio_clip, - motion_lengths=speaker_clip_length, - return_distribution=False, - ) - d_loss = (div_loss(listener_3dmm_out_, listener_3dmm_out) + - div_loss(listener_emotion_out_, listener_emotion_out)) - loss = loss + self.div_p * d_loss - - batch_size = len(speaker_video_clip) - losses.update(loss.data.item(), batch_size) - rec_losses.update(rec_loss.data.item(), batch_size) - rec_emo_losses.update(rec_emo_loss.data.item(), batch_size) - rec_param_losses.update(rec_param_loss.data.item(), batch_size) - kld_losses.update(kld_loss.data.item(), batch_size) - div_losses.update(d_loss.data.item(), batch_size) - eeg_losses.update(eeg_loss.data.item(), batch_size) - eeg_valid_ratios.update(eeg_valid_ratio.data.item(), batch_size) - - loss.backward() - optimizer.step() - return ( - losses.avg, - rec_losses.avg, - rec_emo_losses.avg, - rec_param_losses.avg, - kld_losses.avg, - div_losses.avg, - eeg_losses.avg, - eeg_valid_ratios.avg, - ) - - # Validation - def val(self, model): - losses = AverageMeter() - rec_losses = AverageMeter() - rec_emo_losses = AverageMeter() - rec_param_losses = AverageMeter() - kld_losses = AverageMeter() - eeg_losses = AverageMeter() - eeg_valid_ratios = AverageMeter() - model.eval() - model.reset_window_size(8) - - for batch_idx, batch in enumerate(tqdm(self.val_loader)): - ( - speaker_audio_clip, - speaker_video_clip, - speaker_emotion_clip, - speaker_3dmm_clip, - _, - listener_emotion, - listener_3dmm, - speaker_clip_length, - listener_clip_length, - ) = batch[:9] - listener_eeg_clip = listener_eeg_mask = None - if len(batch) > 9: - listener_eeg_clip, listener_eeg_mask = batch[9:11] - if self.model_cfg.task == 'offline': - (speaker_audio_clip, speaker_video_clip, speaker_emotion_clip, speaker_3dmm_clip, - listener_emotion, listener_3dmm, listener_eeg_clip, listener_eeg_mask) = self.data_resample( - speaker_audio_clip, - speaker_video_clip, - speaker_emotion_clip, - speaker_3dmm_clip, - listener_emotion, - listener_3dmm, - speaker_clip_length, - listener_clip_length, - listener_eeg_clips=listener_eeg_clip, - listener_eeg_masks=listener_eeg_mask, - ) - - with (torch.no_grad()): - listener_3dmm_out, listener_emotion_out, distribution, eeg_outputs = model( - speaker_video_clip, - speaker_audio_clip, - motion_lengths=speaker_clip_length, - speaker_emotion=speaker_emotion_clip, - speaker_3dmm=speaker_3dmm_clip, - listener_eeg_input=listener_eeg_clip, - listener_eeg_mask=listener_eeg_mask, - return_eeg_outputs=True, - ) - loss, rec_loss, rec_emo_loss, rec_param_loss, kld_loss, eeg_loss, eeg_valid_ratio = self.criterion( - listener_emotion, - listener_3dmm, - listener_emotion_out, - listener_3dmm_out, - distribution, - prediction_eeg=eeg_outputs.get("prediction_eeg"), - target_eeg=eeg_outputs.get("target_eeg"), - target_eeg_mask=eeg_outputs.get("target_eeg_mask"), - ) - if self.train_eeg_head_only: - loss = eeg_loss - - batch_size = len(speaker_video_clip) - losses.update(loss.data.item(), batch_size) - rec_losses.update(rec_loss.data.item(), batch_size) - rec_emo_losses.update(rec_emo_loss.data.item(), batch_size) - rec_param_losses.update(rec_param_loss.data.item(), batch_size) - kld_losses.update(kld_loss.data.item(), batch_size) - eeg_losses.update(eeg_loss.data.item(), batch_size) - eeg_valid_ratios.update(eeg_valid_ratio.data.item(), batch_size) - - model.reset_window_size(self.window_size) - return ( - losses.avg, - rec_losses.avg, - rec_emo_losses.avg, - rec_param_losses.avg, - kld_losses.avg, - eeg_losses.avg, - eeg_valid_ratios.avg, - ) - - def pad_to(self, seq: torch.Tensor, length: int) -> torch.Tensor: - L = seq.shape[0] - if L < length: - pad_shape = (length - L, *seq.shape[1:]) - return torch.cat([seq, seq.new_zeros(pad_shape)], dim=0) - return seq - - def test(self): - stage = "test" - data_clamp = self.kwargs.pop("data_clamp") - - model = instantiate(self.model_cfg, _recursive_=False) - checkpoint_path = self.get_ckpt_path(model, runid="resume_runid", best=True) - # checkpoint_path = self.get_ckpt_path(model, runid="resume_runid", epoch=30) - self._load_model_checkpoint(model, checkpoint_path, strict=False) - model.eval() - if self.eval_eeg: - if getattr(model, "eeg_head", None) is None: - raise RuntimeError( - "trainer.eval_eeg=True but configs//model/motion_transvae.yaml has no enabled eeg_head." - ) - eeg_ckpt_path = self.get_ckpt_path( - model.eeg_head, - runid="resume_runid", - best=True, - create_dir=False, - ) - if not os.path.exists(eeg_ckpt_path): - raise FileNotFoundError( - "trainer.eval_eeg=True requires a trained EEGPredictionHead checkpoint. " - f"Missing: {eeg_ckpt_path}" - ) - self._load_model_checkpoint(model.eeg_head, eeg_ckpt_path, strict=True) - model.eval() - - # Instantiate the renderer only when rendering is requested. Hydra changes the - # working directory during evaluation, and renderer init loads large external - # assets that are unnecessary for metrics-only runs. - renderer = None - if self.renderer_cfg.do_render: - renderer = instantiate( - self.renderer_cfg, - device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), - ) - - logger.info("Loading test data module") - test_loader = self.data_module.get_dataloader(stage=stage) - logger.info("Test data module loaded") - - keep_facial_outputs = self.eval_facial_metrics - keep_eeg_outputs = self.eval_eeg and (self.save_results or self.eval_eeg_metrics) - - post_processor = None - post_config_name = self.kwargs.pop("post_config_name") - post_clip_length = self.kwargs.pop("post_clip_length") - if keep_facial_outputs: - logger.info("Loading post processor") - post_processor = Processor( - config_name=post_config_name, - clip_len_test=post_clip_length, - device=self.device, - ) - logger.info("Post processor loaded") - - speaker_emotions_input_all = [] - listener_3dmm_preds_lists_all = [] - listener_emotion_preds_lists_all = [] - listener_3dmm_GTs_all = [] - listener_emotion_GTs_all = [] - listener_eeg_preds_lists_all = [] - listener_eeg_GTs_all = [] - listener_eeg_masks_all = [] - max_seq_len = self.max_seq_len - num_preds = self.num_preds - - for batch_idx, batch in enumerate(tqdm(test_loader)): - ( - speaker_audio_clips, - speaker_video_clips, - speaker_emotion_clips, - speaker_3dmm_clips, - listener_video_clips, - listener_emotions, - listener_3dmms, - speaker_clip_lengths, - listener_clip_lengths, - ) = batch[:9] - listener_eeg_clips = listener_eeg_masks = None - if len(batch) > 9: - listener_eeg_clips, listener_eeg_masks = batch[9:11] - if self.eval_eeg and listener_eeg_clips is None: - raise RuntimeError("trainer.eval_eeg=True but the test dataloader did not return EEG labels.") - if renderer is not None and len(listener_video_clips) == 0: - raise RuntimeError( - "trainer.renderer.do_render=True but the test dataloader did not return listener video. " - "Set data.test_dataset.load_video_l=true when rendering." - ) - listener_video_iter = listener_video_clips if renderer is not None else [None] * len(speaker_audio_clips) - - listener_3dmm_preds = [] if keep_facial_outputs else None - listener_emotion_preds = [] if keep_facial_outputs else None - listener_eeg_preds = [] if keep_eeg_outputs else None - sample_eeg_targets = [] if keep_eeg_outputs else None - sample_eeg_masks = [] if keep_eeg_outputs else None - for i in range(num_preds): - if i == 0: - if keep_facial_outputs: - for em in listener_emotions: - listener_emotion_GTs_all.append([em] if isinstance(em, torch.Tensor) else em) - if self.eval_facial_metrics: - speaker_emotions_input_all.extend(speaker_emotion_clips) - - eeg_clips = listener_eeg_clips if self.eval_eeg else [None] * len(speaker_audio_clips) - eeg_masks = listener_eeg_masks if self.eval_eeg else [None] * len(speaker_audio_clips) - for j, (speaker_audio_clip, speaker_video_clip, speaker_emotion_clip, speaker_3dmm_clip, - listener_video_clip, speaker_clip_length, listener_eeg_clip, listener_eeg_mask) in ( - enumerate(zip(speaker_audio_clips, speaker_video_clips, speaker_emotion_clips, - speaker_3dmm_clips, listener_video_iter, speaker_clip_lengths, - eeg_clips, eeg_masks))): - - speaker_audio_clip_list = [] - speaker_video_clip_list = [] - speaker_emotion_clip_list = [] - speaker_3dmm_clip_list = [] - listener_eeg_clip_list = [] - listener_eeg_mask_list = [] - motion_length_list = [] - speaker_clip_length_int = int( - speaker_clip_length.item() if torch.is_tensor(speaker_clip_length) else speaker_clip_length - ) - - # split into sub-clips - for k in range(math.ceil(speaker_clip_length_int / max_seq_len)): - start_idx = k * max_seq_len - end_idx = min((k + 1) * max_seq_len, speaker_clip_length_int) - motion_length = end_idx - start_idx - motion_length_list.append(motion_length) - speaker_audio_clip_list.append(speaker_audio_clip[start_idx:end_idx]) - speaker_video_clip_list.append(speaker_video_clip[start_idx:end_idx]) - speaker_emotion_clip_list.append(speaker_emotion_clip[start_idx:end_idx]) - speaker_3dmm_clip_list.append(speaker_3dmm_clip[start_idx:end_idx]) - if self.eval_eeg: - listener_eeg_clip_list.append(listener_eeg_clip[start_idx:end_idx]) - listener_eeg_mask_list.append(listener_eeg_mask[start_idx:end_idx]) - - if self.model_cfg.task == 'online': - speaker_audio_clip_list[-1] = self.pad_to(speaker_audio_clip_list[-1], max_seq_len) - speaker_video_clip_list[-1] = self.pad_to(speaker_video_clip_list[-1], max_seq_len) - speaker_emotion_clip_list[-1] = self.pad_to(speaker_emotion_clip_list[-1], max_seq_len) - speaker_3dmm_clip_list[-1] = self.pad_to(speaker_3dmm_clip_list[-1], max_seq_len) - if self.eval_eeg: - listener_eeg_clip_list[-1] = self.pad_to(listener_eeg_clip_list[-1], max_seq_len) - listener_eeg_mask_list[-1] = self.pad_to(listener_eeg_mask_list[-1], max_seq_len) - speaker_audio_clip_inputs = speaker_audio_clip_list # List: [tensor([l, d_audio]), ...] - speaker_video_clip_inputs = speaker_video_clip_list # List: [tensor([l, 3, 224, 224]), ...] - speaker_emotion_clip_inputs = speaker_emotion_clip_list - speaker_3dmm_clip_inputs = speaker_3dmm_clip_list - listener_eeg_clip_inputs = listener_eeg_clip_list if self.eval_eeg else None - listener_eeg_mask_inputs = listener_eeg_mask_list if self.eval_eeg else None - - chunk_3dmm_outputs = [] - chunk_emotion_outputs = [] - chunk_eeg_outputs = [] - chunk_eeg_targets = [] - chunk_eeg_masks = [] - keep_render_outputs = renderer is not None - num_chunks = len(speaker_audio_clip_inputs) - for chunk_start in range(0, num_chunks, self.eval_clip_batch_size): - chunk_end = min(chunk_start + self.eval_clip_batch_size, num_chunks) - with torch.inference_mode(): - mb_3dmm_out, mb_emotion_out, _, mb_eeg_outputs = model( - speaker_video_clip_inputs[chunk_start:chunk_end], - speaker_audio_clip_inputs[chunk_start:chunk_end], - motion_lengths=torch.tensor(motion_length_list[chunk_start:chunk_end]), - speaker_emotion=speaker_emotion_clip_inputs[chunk_start:chunk_end], - speaker_3dmm=speaker_3dmm_clip_inputs[chunk_start:chunk_end], - listener_eeg_input=listener_eeg_clip_inputs[chunk_start:chunk_end] - if self.eval_eeg else None, - listener_eeg_mask=listener_eeg_mask_inputs[chunk_start:chunk_end] - if self.eval_eeg else None, - return_eeg_outputs=True, - return_distribution=False, - ) - - if keep_facial_outputs or keep_render_outputs: - chunk_3dmm_outputs.extend(output.detach().cpu() for output in mb_3dmm_out) - chunk_emotion_outputs.extend(output.detach().cpu() for output in mb_emotion_out) - - if self.eval_eeg: - if "prediction_eeg" not in mb_eeg_outputs: - raise RuntimeError( - "trainer.eval_eeg=True but the model did not return prediction_eeg." - ) - if "target_eeg" not in mb_eeg_outputs: - raise RuntimeError( - "trainer.eval_eeg=True but the model did not return target_eeg." - ) - if keep_eeg_outputs: - chunk_eeg_outputs.append(mb_eeg_outputs["prediction_eeg"].detach().cpu()) - chunk_eeg_targets.append(mb_eeg_outputs["target_eeg"].detach().cpu()) - chunk_eeg_masks.append(mb_eeg_outputs["target_eeg_mask"].detach().cpu()) - - del mb_3dmm_out, mb_emotion_out, mb_eeg_outputs - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - listener_3dmm_out = listener_emotion_out = None - if keep_facial_outputs or keep_render_outputs: - listener_3dmm_out = torch.cat(chunk_3dmm_outputs, dim=0)[:speaker_clip_length_int] - listener_3dmm_out = listener_3dmm_out.unsqueeze(0) - listener_emotion_out = torch.cat(chunk_emotion_outputs, dim=0)[:speaker_clip_length_int] - listener_emotion_out = listener_emotion_out.unsqueeze(0) - if data_clamp: - listener_emotion_out[:, :, :15] = torch.round(listener_emotion_out[:, :, :15]) - - if keep_eeg_outputs: - listener_eeg_out = torch.cat(chunk_eeg_outputs, dim=0).unsqueeze(0) - listener_eeg_target = torch.cat(chunk_eeg_targets, dim=0).unsqueeze(0) - listener_eeg_mask_out = torch.cat(chunk_eeg_masks, dim=0).unsqueeze(0) - - if renderer is not None and i == 0: # (batch_idx % 20) == 0 - listener_video_clip = listener_video_clip[0].to(self.device) - val_path = os.path.join('results_videos', 'test') - os.makedirs(val_path, exist_ok=True) - - perm = torch.randperm(listener_video_clip.shape[0]) - listener_references = listener_video_clip[perm[0]] - assert len(listener_references.shape) == 3, \ - "listener_references.shape should be (3, 224, 224)" - renderer.rendering(val_path, - f"batch{str(batch_idx + 1)}", - listener_3dmm_out.to(self.device), - speaker_video_clip, - listener_references, - listener_video_clip) - - if i == 0: - if keep_facial_outputs: - listener_3dmm_preds.append(listener_3dmm_out) - listener_emotion_preds.append(listener_emotion_out) - if keep_eeg_outputs: - listener_eeg_preds.append(listener_eeg_out) - sample_eeg_targets.append(listener_eeg_target) - sample_eeg_masks.append(listener_eeg_mask_out) - else: - if keep_facial_outputs: - listener_3dmm_preds[j] = torch.cat( - (listener_3dmm_preds[j], listener_3dmm_out), dim=0) - listener_emotion_preds[j] = torch.cat( - (listener_emotion_preds[j], listener_emotion_out), dim=0) - if keep_eeg_outputs: - listener_eeg_preds[j] = torch.cat( - (listener_eeg_preds[j], listener_eeg_out), dim=0) - - # listener_3dmm_preds: (num_preds, l, ...) - if keep_facial_outputs: - listener_3dmm_preds_lists_all.extend(listener_3dmm_preds) - # listener_emotion_preds: (num_preds, l, ...) - listener_emotion_preds_lists_all.extend(listener_emotion_preds) - if keep_eeg_outputs: - listener_eeg_preds_lists_all.extend(listener_eeg_preds) - listener_eeg_GTs_all.extend(sample_eeg_targets) - listener_eeg_masks_all.extend(sample_eeg_masks) - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - # listener_emotion_preds_lists_all - # List: 750 [Tensor([num_preds, l, 25]), Tensor([num_preds, l', 25]), ...] - # listener_emotion_GTs_all - # List: 750 [List: [(l', 25), (l'', 25), ...], List: [(l''', 25), (l'''', 25)], ...] - if keep_facial_outputs and len(listener_emotion_preds_lists_all): - listener_emotion_GTs_all = post_processor.forward( - prediction_list=listener_emotion_preds_lists_all, - target_list=listener_emotion_GTs_all,) - # listener_emotion_GTs_all - # List: 750 [Tensor([num_preds, l, 25]), Tensor([num_preds, l', 25]), ...] - - if self.save_results: - try: - result_dict = {} - if keep_facial_outputs: - result_dict.update({'GT': listener_emotion_GTs_all, 'PRED': listener_emotion_preds_lists_all}) - if keep_eeg_outputs: - result_dict.update({ - 'GT_EEG': listener_eeg_GTs_all, - 'PRED_EEG': listener_eeg_preds_lists_all, - 'EEG_MASK': listener_eeg_masks_all, - }) - torch.save(result_dict, 'results.pt') - print("Successfully saved Tensor List") - except Exception: - print("Failed to save Tensor List") - - results = {} - if self.eval_facial_metrics: - results.update(compute_metrics( - speaker_emotions_input_all, - listener_emotion_preds_lists_all, - listener_emotion_GTs_all, - threads=self.metric_threads, - )) - if self.eval_eeg and self.eval_eeg_metrics: - results.update(compute_eeg_metrics( - listener_eeg_preds_lists_all, - listener_eeg_GTs_all, - listener_eeg_masks_all, - )) - logger.info(results) diff --git a/personalised/code/trainer/perfrdiff_rewrite_weight.py b/personalised/code/trainer/perfrdiff_rewrite_weight.py deleted file mode 100644 index 18cfecb35cfc7f30d1d9a2e1ed226d3d99f10713..0000000000000000000000000000000000000000 --- a/personalised/code/trainer/perfrdiff_rewrite_weight.py +++ /dev/null @@ -1,1049 +0,0 @@ -import logging -import math -import os -from functools import partial -from pathlib import Path - -import hydra -import torch -from einops import rearrange -from hydra.utils import instantiate -from omegaconf import DictConfig, OmegaConf -from torch import optim -from torch.utils.tensorboard import SummaryWriter -from tqdm import tqdm - -from framework.modules.post_processor import Processor -from framework.perfrdiff_rewrite_weight import losses as rewrite_losses -from framework.perfrdiff_rewrite_weight.modifier.network import MainNetUnified -from framework.utils.compute_metrics import compute_eeg_metrics, compute_metrics -from framework.utils.util import from_pretrained_checkpoint -from utils.util import AverageMeter - -logger = logging.getLogger(__name__) - - -class Trainer: - def __init__( - self, - resumed_training: bool = False, - generic: DictConfig = None, - model: DictConfig = None, - criterion: DictConfig = None, - person_specific: DictConfig = None, - main_model: DictConfig = None, - pretrained: DictConfig = None, - batch_size: int = 4, - post_config_name: str = "configs/shared/model/emotion_autoencoder.yaml", - post_clip_length: int = 1000, - data_clamp: bool = True, - num_eval_preds: int = 10, - eval_clip_batch_size: int = 8, - parallel_eval_preds: bool = True, - **kwargs, - ): - super().__init__() - self.resumed_training = resumed_training - self.trainer_cfg = generic - self.model_cfg = model - self.criterion_cfg = criterion - self.person_specific_cfg = person_specific - self.main_model_cfg = main_model - self.pretrained_cfg = pretrained - self.batch_size = batch_size - self.post_config_name = post_config_name - self.post_clip_length = post_clip_length - self.data_clamp = data_clamp - self.num_eval_preds = num_eval_preds - self.eval_clip_batch_size = eval_clip_batch_size - self.parallel_eval_preds = parallel_eval_preds - self.kwargs = kwargs - self.task = kwargs.get("task", "online") - self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - self.train_eeg = self._as_bool(self.trainer_cfg.get("train_eeg", False)) - self.train_eeg_head_only = self._as_bool(self.trainer_cfg.get("train_eeg_head_only", False)) - self.train_eeg = self.train_eeg or self.train_eeg_head_only - self.eval_eeg = self._as_bool(self.trainer_cfg.get("eval_eeg", False)) - self.skip_modifier = self._as_bool(self.trainer_cfg.get("skip_modifier", False)) - - @staticmethod - def _as_bool(value): - if isinstance(value, str): - return value.lower() in {"1", "true", "yes", "y"} - return bool(value) - - @staticmethod - def _resolve_checkpoint_path(path): - if path is None or str(path).strip() == "": - return None - path = str(path) - if os.path.isabs(path): - return path - return hydra.utils.to_absolute_path(path) - - def set_data_module(self, data_module): - self.data_module = data_module - - def _ensure_eeg_data_enabled(self, stage): - if stage == "fit" and self.train_eeg: - if hasattr(self.data_module, "train_set_cfg"): - self.data_module.train_set_cfg.load_eeg_l = True - if hasattr(self.data_module, "val_set_cfg"): - self.data_module.val_set_cfg.load_eeg_l = True - if stage == "test" and self.eval_eeg and hasattr(self.data_module, "test_set_cfg"): - self.data_module.test_set_cfg.load_eeg_l = True - - def _rewrite_cfg(self): - return OmegaConf.create( - { - "main_model": OmegaConf.to_container(self.main_model_cfg, resolve=True), - "person_specific": OmegaConf.to_container(self.person_specific_cfg, resolve=True), - } - ) - - def _build_diffusion(self, stage): - model_cfg = self.model_cfg - if stage == "test" and self.parallel_eval_preds and self.num_eval_preds > 1: - model_cfg = OmegaConf.create(OmegaConf.to_container(self.model_cfg, resolve=True)) - if model_cfg.diff_model.get("diffusion_prior") is not None: - model_cfg.diff_model.diffusion_prior.scheduler.num_preds = self.num_eval_preds - model_cfg.diff_model.diffusion_decoder.scheduler.num_preds = self.num_eval_preds - - model = instantiate( - model_cfg.diff_model, - stage=stage, - resumed_training=False, - auto_load_ckpt=False, - latent_embedder=model_cfg.latent_embedder - if hasattr(model_cfg, "latent_embedder") else None, - audio_encoder=model_cfg.audio_encoder - if hasattr(model_cfg, "audio_encoder") else None, - **self.kwargs, - _recursive_=False, - ) - model.to(self.device) - self._load_pretrained_diffusion(model) - return model - - def _load_pretrained_diffusion(self, model): - if self.pretrained_cfg is None: - raise ValueError("Missing trainer.pretrained configuration for rewrite-weight diffusion.") - - def load_required(path, module, label): - checkpoint_path = hydra.utils.to_absolute_path(path) - if not os.path.isfile(checkpoint_path): - raise FileNotFoundError( - f"Missing pretrained {label} checkpoint: {checkpoint_path}. " - "Please place the pretrained diffusion weights under pretrained_models/ " - "or override trainer.pretrained.*." - ) - from_pretrained_checkpoint(checkpoint_path, module, self.device) - - if getattr(model, "diffusion_prior", None) is not None: - load_required(self.pretrained_cfg.diffusion_prior, model.diffusion_prior.model, "DiffusionPriorNetwork") - load_required(self.pretrained_cfg.diffusion_decoder, model.diffusion_decoder.model, "TransformerDenoiser") - - def _build_model(self, stage): - diffusion = self._build_diffusion(stage) - model = MainNetUnified(self._rewrite_cfg(), diffusion, self.device) - model.to(self.device) - return model - - def _build_criterion(self): - return partial(getattr(rewrite_losses, self.criterion_cfg.type), **self.criterion_cfg.args) - - def _build_optimizer(self, model): - cfg = self.main_model_cfg.optimizer_hypernet.args - if self.train_eeg_head_only: - model.freeze_except_eeg_head() - params = list(model.eeg_head().parameters()) - trainable_names = [name for name, parameter in model.named_parameters() if parameter.requires_grad] - logger.info("EEG head-only trainable parameter tensors: %s", trainable_names) - print(f"EEG head-only trainable parameter tensors: {trainable_names}") - else: - params = list(model.modifier_parameters(include_eeg_head=self.train_eeg)) - return optim.SGD( - params=params, - lr=cfg.lr, - momentum=cfg.momentum, - weight_decay=cfg.weight_decay, - ) - - def _personal_condition_mode(self): - mode = self.main_model_cfg.args.get("personal_condition_mode", "3dmm_personality") - if mode == "history_personality": - return "3dmm_personality" - return mode - - @staticmethod - def _infer_checkpoint_personal_condition_mode(checkpoint): - mode = checkpoint.get("personal_condition_mode") - if mode == "history_personality": - return "3dmm_personality" - if mode is not None: - return mode - - state_dict = checkpoint.get("state_dict", {}) - if any(key.startswith("personality_fusion.") for key in state_dict): - return "3dmm_personality" - if any(key.startswith("personality_encoder.") for key in state_dict): - return "personality_only" - return "3dmm_only" - - def _validate_checkpoint_personal_condition_mode(self, checkpoint, checkpoint_path): - checkpoint_mode = self._infer_checkpoint_personal_condition_mode(checkpoint) - current_mode = self._personal_condition_mode() - if checkpoint_mode != current_mode: - raise ValueError( - "ModifierNetwork checkpoint personal_condition_mode mismatch. " - f"Checkpoint: {checkpoint_mode}; current config: {current_mode}; path: {checkpoint_path}. " - "Use a checkpoint trained with the same mode, or set " - f"trainer.main_model.args.personal_condition_mode={checkpoint_mode} for evaluation." - ) - - @staticmethod - def _checkpoint_has_eeg_head(checkpoint): - state_dict = checkpoint.get("state_dict", {}) - return any(key.startswith("eeg_head.") for key in state_dict) - - def _load_eeg_head_checkpoint(self, model, path): - checkpoint_path = self._resolve_checkpoint_path(path) - if checkpoint_path is None: - return - if not os.path.isfile(checkpoint_path): - raise FileNotFoundError(f"Missing EEG head checkpoint: {checkpoint_path}") - checkpoint = torch.load(checkpoint_path, map_location="cpu") - state_dict = checkpoint.get("state_dict", checkpoint) - if any(key.startswith("eeg_head.") for key in state_dict): - state_dict = { - key[len("eeg_head."):]: value - for key, value in state_dict.items() - if key.startswith("eeg_head.") - } - if not model.has_eeg_head(): - raise RuntimeError("Cannot load EEG head checkpoint because main_net.eeg_head is disabled.") - model.eeg_head().load_state_dict(state_dict) - model.to(self.device) - logger.info("Loaded EEG head checkpoint: %s", checkpoint_path) - - def _load_modifier_checkpoint_file(self, model, path, require_eeg_head=False): - checkpoint_path = self._resolve_checkpoint_path(path) - if checkpoint_path is None: - raise ValueError("trainer.pretrained.modifier_checkpoint is required for EEG head-only training.") - if not os.path.isfile(checkpoint_path): - raise FileNotFoundError(f"Missing ModifierNetwork checkpoint: {checkpoint_path}") - checkpoint = torch.load(checkpoint_path, map_location="cpu") - self._validate_checkpoint_personal_condition_mode(checkpoint, checkpoint_path) - if require_eeg_head and not self._checkpoint_has_eeg_head(checkpoint): - raise RuntimeError(f"ModifierNetwork checkpoint has no trained EEG head: {checkpoint_path}") - model.load_modifier_state_dict(checkpoint["state_dict"]) - model.to(self.device) - logger.info("Loaded pretrained ModifierNetwork checkpoint: %s", checkpoint_path) - - def _modifier_dir(self, run_key="current_runid"): - ckpt_root = Path(hydra.utils.to_absolute_path(self.kwargs.get("ckpt_dir"))) - run_id = self.kwargs.get(run_key) or self.kwargs.get("current_runid") - ckpt_dir = ckpt_root / str(run_id) / "ModifierNetwork" - ckpt_dir.mkdir(parents=True, exist_ok=True) - return ckpt_dir - - def _save_checkpoint(self, model, optimizer, epoch=None, best=False, last=False, - save_epoch=False, best_loss=float("inf")): - ckpt_dir = self._modifier_dir("current_runid") - def save_modifier_checkpoint(path): - checkpoint = { - "epoch": epoch if epoch is not None else None, - "best_loss": best_loss if best_loss is not None else None, - "personal_condition_mode": self._personal_condition_mode(), - "state_dict": model.modifier_state_dict(include_eeg_head=self.train_eeg), - "train_eeg": self.train_eeg, - "optimizer": optimizer.state_dict() if optimizer is not None else None, - } - torch.save(checkpoint, str(path)) - - if save_epoch and epoch is not None: - save_modifier_checkpoint(ckpt_dir / f"checkpoint_{epoch}.pth") - if best: - save_modifier_checkpoint(ckpt_dir / "checkpoint_best.pth") - if last: - save_modifier_checkpoint(ckpt_dir / "checkpoint_last.pth") - - def _load_modifier_checkpoint(self, model, optimizer=None, run_key="resume_runid", - names=None, require_eeg_head=False): - ckpt_dir = self._modifier_dir(run_key) - names = names or ["checkpoint_best.pth", "checkpoint_last.pth"] - for name in names: - checkpoint_path = ckpt_dir / name - if checkpoint_path.is_file(): - checkpoint = torch.load(str(checkpoint_path), map_location="cpu") - self._validate_checkpoint_personal_condition_mode(checkpoint, checkpoint_path) - if require_eeg_head and not self._checkpoint_has_eeg_head(checkpoint): - raise RuntimeError(f"ModifierNetwork checkpoint has no trained EEG head: {checkpoint_path}") - model.load_modifier_state_dict(checkpoint["state_dict"]) - if optimizer is not None and checkpoint.get("optimizer") is not None: - try: - optimizer.load_state_dict(checkpoint["optimizer"]) - for state in optimizer.state.values(): - for key, value in state.items(): - if torch.is_tensor(value): - state[key] = value.to(self.device) - except ValueError: - logger.warning( - "Skip optimizer state from %s because the modifier parameter set changed.", - checkpoint_path, - ) - model.to(self.device) - logger.info("Loaded ModifierNetwork checkpoint: %s", checkpoint_path) - return checkpoint.get("best_loss", float("inf")), checkpoint.get("epoch", 0) - raise FileNotFoundError(f"No ModifierNetwork checkpoint found in {ckpt_dir}; tried {names}") - - def _resample_train_batch(self, speaker_audio, speaker_emotion, speaker_3dmm, listener_emotion, - listener_eeg=None, listener_eeg_mask=None): - window_size = self.trainer_cfg.window_size - clip_length = self.trainer_cfg.clip_length - s_window_size = self.trainer_cfg.s_ratio * window_size - has_eeg = listener_eeg is not None and listener_eeg.numel() > 0 - - if self.task == "offline": - motion_lengths = torch.full( - (speaker_audio.shape[0],), - min(clip_length, speaker_audio.shape[1]), - dtype=torch.long, - ) - eeg_target = listener_eeg[:, motion_lengths[0] - 1] if has_eeg else None - eeg_mask = listener_eeg_mask[:, motion_lengths[0] - 1] if has_eeg else None - return speaker_audio, speaker_emotion, speaker_3dmm, listener_emotion, None, motion_lengths, eeg_target, eeg_mask - - if self.task != "online": - raise ValueError(f"Unknown task type: {self.task}") - - sampled_audio = [] - sampled_emotion = [] - sampled_3dmm = [] - sampled_listener = [] - sampled_past_listener = [] - sampled_eeg = [] - sampled_eeg_mask = [] - for idx in range(speaker_audio.shape[0]): - seq_length = min(clip_length, speaker_audio.shape[1], listener_emotion.shape[1]) - max_start = max(seq_length - s_window_size, 0) - cp = torch.randint(0, max_start + 1, (1,)).item() if max_start > 0 else 0 - du = cp + s_window_size - - sampled_audio.append(speaker_audio[idx, cp:du]) - sampled_emotion.append(speaker_emotion[idx, cp:du]) - sampled_3dmm.append(speaker_3dmm[idx, cp:du]) - sampled_past_listener.append(listener_emotion[idx, du - 2 * window_size:du - window_size]) - sampled_listener.append(listener_emotion[idx, du - window_size:du]) - if has_eeg: - sampled_eeg.append(listener_eeg[idx, du - 1]) - sampled_eeg_mask.append(listener_eeg_mask[idx, du - 1]) - - eeg_target = torch.stack(sampled_eeg, dim=0) if has_eeg else None - eeg_mask = torch.stack(sampled_eeg_mask, dim=0) if has_eeg else None - return ( - torch.stack(sampled_audio, dim=0), - torch.stack(sampled_emotion, dim=0), - torch.stack(sampled_3dmm, dim=0), - torch.stack(sampled_listener, dim=0), - torch.stack(sampled_past_listener, dim=0), - None, - eeg_target, - eeg_mask, - ) - - def _split_outputs(self, outputs): - if isinstance(outputs, dict) and "output_prior" in outputs: - return outputs["output_prior"], outputs["output_decoder"] - if isinstance(outputs, (list, tuple)) and len(outputs) == 2: - return outputs - raise ValueError("Rewrite-weight training expects diffusion output with prior and decoder branches.") - - def _single_loss(self, model, criterion, speaker_audio, speaker_emotion, speaker_3dmm, - listener_emotion, past_listener_emotion, motion_length, - personal_3dmm, listener_personality, listener_eeg, listener_eeg_mask, idx): - input_dict = { - "speaker_audio_input": speaker_audio[idx:idx + 1], - "speaker_emotion_input": speaker_emotion[idx:idx + 1], - "speaker_3dmm_input": speaker_3dmm[idx:idx + 1], - "listener_emotion_input": listener_emotion[idx:idx + 1], - "past_listener_emotion": past_listener_emotion[idx:idx + 1] if past_listener_emotion is not None else None, - "motion_length": motion_length[idx:idx + 1].to(self.device) if motion_length is not None else None, - "listener_eeg_input": listener_eeg[idx:idx + 1] if listener_eeg is not None else None, - "listener_eeg_mask": listener_eeg_mask[idx:idx + 1] if listener_eeg_mask is not None else None, - } - personal_3dmm_input = personal_3dmm[idx:idx + 1] if personal_3dmm.numel() > 0 else None - outputs, regular_loss = model( - x=input_dict, - p=personal_3dmm_input, - personality=listener_personality[idx:idx + 1], - ) - output_prior, output_decoder = self._split_outputs(outputs) - loss_dict = criterion(output_prior, output_decoder) - loss = loss_dict["loss_eeg"] if self.train_eeg_head_only else loss_dict["loss"] + regular_loss - if self.train_eeg_head_only and torch.is_grad_enabled() and not loss.requires_grad: - raise RuntimeError("loss_eeg has no gradient. Check EEG labels and prediction_eeg.") - return loss, loss_dict, regular_loss - - def fit(self): - if self.train_eeg_head_only and self.resumed_training: - raise ValueError("train_eeg_head_only=True should be launched with resume=false.") - - logger.info("Loading rewrite-weight data module") - self._ensure_eeg_data_enabled(stage="fit") - train_loader, val_loader = self.data_module.get_dataloader(stage="fit") - logger.info("Data module loaded") - - model = self._build_model(stage="fit") - if self.train_eeg_head_only: - self._load_modifier_checkpoint_file( - model, - self.pretrained_cfg.get("modifier_checkpoint", ""), - require_eeg_head=False, - ) - elif not self.resumed_training: - # (online recipe) warm-start the modifier from an OFFLINE personalized - # checkpoint so online fine-tuning starts from a competent listener model - # (analogous to warm-starting generic online from generic offline). Keys - # absent in the source (e.g. person_coarse) keep their fresh init. - ws_path = self._resolve_checkpoint_path( - self.pretrained_cfg.get("modifier_warmstart", "") if self.pretrained_cfg else "") - if ws_path and os.path.isfile(ws_path): - ckpt = torch.load(ws_path, map_location="cpu") - self._validate_checkpoint_personal_condition_mode(ckpt, ws_path) - model.load_modifier_state_dict(ckpt["state_dict"]) - model.to(self.device) - logger.info("Warm-started modifier from offline checkpoint: %s", ws_path) - self._load_eeg_head_checkpoint(model, self.pretrained_cfg.get("eeg_head_checkpoint", "")) - optimizer = self._build_optimizer(model) - criterion = self._build_criterion() - writer = SummaryWriter(self.trainer_cfg.tb_dir) - - best_loss = float("inf") - start_epoch = self.trainer_cfg.start_epoch - if self.resumed_training: - best_loss, start_epoch = self._load_modifier_checkpoint( - model, - optimizer=optimizer, - run_key="resume_runid", - names=["checkpoint_last.pth"], - ) - logger.info("Resume rewrite-weight training from epoch %s", start_epoch) - - selected_loss_name = "loss_eeg" if self.train_eeg_head_only else "loss" - - for epoch in range(start_epoch, self.trainer_cfg.epochs): - train_loss, train_prior, train_decoder, train_eeg_loss, train_eeg_valid = self._run_epoch( - model, train_loader, criterion, optimizer, writer, epoch, train=True - ) - logger.info( - "Epoch: %s train_%s: %.5f prior_loss: %.5f decoder_loss: %.5f " - "eeg_loss: %.5f eeg_valid_ratio: %.5f", - epoch + 1, selected_loss_name, train_loss, train_prior, train_decoder, - train_eeg_loss, train_eeg_valid, - ) - - val_loss = train_loss - val_eeg_loss = train_eeg_loss - if (epoch + 1) % self.trainer_cfg.val_period == 0: - val_loss, val_prior, val_decoder, val_eeg_loss, val_eeg_valid = self._run_epoch( - model, val_loader, criterion, None, writer, epoch, train=False - ) - logger.info( - "Epoch: %s val_%s: %.5f prior_loss: %.5f decoder_loss: %.5f " - "eeg_loss: %.5f eeg_valid_ratio: %.5f", - epoch + 1, selected_loss_name, val_loss, val_prior, val_decoder, - val_eeg_loss, val_eeg_valid, - ) - selected_val_loss = val_eeg_loss if self.train_eeg_head_only else val_loss - if selected_val_loss < best_loss: - best_loss = selected_val_loss - self._save_checkpoint( - model, optimizer, epoch + 1, best=True, save_epoch=True, best_loss=best_loss, - ) - - if (epoch + 1) % self.trainer_cfg.save_period == 0: - self._save_checkpoint(model, optimizer, epoch + 1, save_epoch=True, best_loss=best_loss) - self._save_checkpoint(model, optimizer, epoch + 1, last=True, best_loss=best_loss) - - writer.close() - - def _run_epoch(self, model, data_loader, criterion, optimizer, writer, epoch, train=True): - whole_losses = AverageMeter() - prior_losses = AverageMeter() - decoder_losses = AverageMeter() - regular_losses = AverageMeter() - eeg_losses = AverageMeter() - eeg_valid_ratios = AverageMeter() - - if train and self.train_eeg_head_only: - model.set_eeg_head_train_mode() - else: - model.train(train) - if model.person_encoder is not None: - model.person_encoder.eval() - iterator = tqdm(data_loader) - - for batch_idx, batch in enumerate(iterator): - if len(batch) == 12: - ( - speaker_audio, - _, - speaker_emotion, - speaker_3dmm, - _, - listener_emotion, - _listener_3dmm, - personal_3dmm, - listener_personality, - listener_eeg, - listener_eeg_mask, - _, - ) = batch - elif len(batch) == 10: - ( - speaker_audio, - _, - speaker_emotion, - speaker_3dmm, - _, - listener_emotion, - _listener_3dmm, - personal_3dmm, - listener_personality, - _, - ) = batch - listener_eeg = listener_eeg_mask = None - else: - ( - speaker_audio, - _, - speaker_emotion, - speaker_3dmm, - _, - listener_emotion, - _listener_3dmm, - personal_3dmm, - _, - ) = batch - listener_personality = speaker_audio.new_zeros((speaker_audio.shape[0], 0)) - listener_eeg = listener_eeg_mask = None - speaker_audio = speaker_audio.to(self.device) - speaker_emotion = speaker_emotion.to(self.device) - speaker_3dmm = speaker_3dmm.to(self.device) - listener_emotion = listener_emotion.to(self.device) - personal_3dmm = personal_3dmm.to(self.device) - listener_personality = listener_personality.to(self.device) - listener_eeg = listener_eeg.to(self.device) if listener_eeg is not None else None - listener_eeg_mask = listener_eeg_mask.to(self.device) if listener_eeg_mask is not None else None - - (speaker_audio, - speaker_emotion, - speaker_3dmm, - listener_emotion, - past_listener_emotion, - motion_length, - listener_eeg, - listener_eeg_mask) = self._resample_train_batch( - speaker_audio, speaker_emotion, speaker_3dmm, listener_emotion, - listener_eeg=listener_eeg, listener_eeg_mask=listener_eeg_mask, - ) - speaker_audio = speaker_audio.to(self.device) - speaker_emotion = speaker_emotion.to(self.device) - speaker_3dmm = speaker_3dmm.to(self.device) - listener_emotion = listener_emotion.to(self.device) - past_listener_emotion = past_listener_emotion.to(self.device) if past_listener_emotion is not None else None - motion_length = motion_length.to(self.device) if motion_length is not None else None - listener_eeg = listener_eeg.to(self.device) if listener_eeg is not None else None - listener_eeg_mask = listener_eeg_mask.to(self.device) if listener_eeg_mask is not None else None - - batch_size = speaker_audio.shape[0] - if train: - optimizer.zero_grad(set_to_none=True) - - batch_loss_value = 0.0 - batch_prior_value = 0.0 - batch_decoder_value = 0.0 - batch_regular_value = 0.0 - batch_eeg_value = 0.0 - batch_eeg_valid_value = 0.0 - context = torch.enable_grad() if train else torch.no_grad() - with context: - for idx in range(batch_size): - loss, loss_dict, regular_loss = self._single_loss( - model, - criterion, - speaker_audio, - speaker_emotion, - speaker_3dmm, - listener_emotion, - past_listener_emotion, - motion_length, - personal_3dmm, - listener_personality, - listener_eeg, - listener_eeg_mask, - idx, - ) - if train: - (loss / batch_size).backward() - batch_loss_value += loss.detach().item() - batch_prior_value += loss_dict["encoded"].detach().item() - batch_decoder_value += loss_dict["decoded"].detach().item() - batch_regular_value += regular_loss.detach().item() - batch_eeg_value += loss_dict["loss_eeg"].detach().item() - batch_eeg_valid_value += loss_dict["eeg_valid_ratio"].detach().item() - - if train: - if self.trainer_cfg.clip_grad: - trainable_params = [parameter for parameter in model.parameters() if parameter.requires_grad] - torch.nn.utils.clip_grad_norm_(trainable_params, 1.0) - optimizer.step() - - batch_loss_value /= batch_size - batch_prior_value /= batch_size - batch_decoder_value /= batch_size - batch_regular_value /= batch_size - batch_eeg_value /= batch_size - batch_eeg_valid_value /= batch_size - whole_losses.update(batch_loss_value, batch_size) - prior_losses.update(batch_prior_value, batch_size) - decoder_losses.update(batch_decoder_value, batch_size) - regular_losses.update(batch_regular_value, batch_size) - eeg_losses.update(batch_eeg_value, batch_size) - eeg_valid_ratios.update(batch_eeg_valid_value, batch_size) - - iteration = batch_idx + len(data_loader) * epoch - if writer is not None: - prefix = "Train" if train else "Val" - writer.add_scalar(f"{prefix}/loss", batch_loss_value, iteration) - writer.add_scalar(f"{prefix}/loss_prior", batch_prior_value, iteration) - writer.add_scalar(f"{prefix}/loss_decoder", batch_decoder_value, iteration) - writer.add_scalar(f"{prefix}/regular_loss", batch_regular_value, iteration) - writer.add_scalar(f"{prefix}/loss_eeg", batch_eeg_value, iteration) - writer.add_scalar(f"{prefix}/eeg_valid_ratio", batch_eeg_valid_value, iteration) - - return whole_losses.avg, prior_losses.avg, decoder_losses.avg, eeg_losses.avg, eeg_valid_ratios.avg - - def _build_test_windows(self, speaker_audio, speaker_emotion, speaker_3dmm, length): - clip_len = self.trainer_cfg.clip_length - window_size = self.trainer_cfg.window_size - s_window_size = self.trainer_cfg.s_ratio * window_size - length = int(length.item() if torch.is_tensor(length) else length) - - if self.task == "offline": - num_windows = max(math.ceil(length / clip_len), 1) - pad_len = num_windows * clip_len - length - motion_lengths = torch.tensor( - [clip_len] * (num_windows - 1) + [length - clip_len * (num_windows - 1)], - dtype=torch.long, - ) - - def pad_and_rearrange(clip): - if pad_len > 0: - clip = torch.cat((clip, clip.new_zeros((pad_len, clip.shape[-1]))), dim=0) - return rearrange(clip, "(b l) d -> b l d", b=num_windows) - - return ( - pad_and_rearrange(speaker_audio[:length]), - pad_and_rearrange(speaker_emotion[:length]), - pad_and_rearrange(speaker_3dmm[:length]), - motion_lengths, - ) - - if self.task != "online": - raise ValueError(f"Unknown task type: {self.task}") - - num_windows = max(math.ceil(length / window_size), 1) - - def pad_online(clip): - return torch.cat( - ( - clip.new_zeros((s_window_size - window_size, clip.shape[-1])), - clip[:length], - clip.new_zeros((num_windows * window_size - length, clip.shape[-1])), - ), - dim=0, - ) - - speaker_audio = pad_online(speaker_audio) - speaker_emotion = pad_online(speaker_emotion) - speaker_3dmm = pad_online(speaker_3dmm) - - audio_windows = [] - emotion_windows = [] - coeff_windows = [] - motion_lengths = [] - for idx in range(num_windows): - start = idx * window_size - audio_windows.append(speaker_audio[start:start + s_window_size]) - emotion_windows.append(speaker_emotion[start:start + s_window_size]) - coeff_windows.append(speaker_3dmm[start:start + s_window_size]) - motion_lengths.append(window_size if idx < num_windows - 1 else length - idx * window_size) - - return ( - torch.stack(audio_windows, dim=0), - torch.stack(emotion_windows, dim=0), - torch.stack(coeff_windows, dim=0), - torch.tensor(motion_lengths, dtype=torch.long), - ) - - def _apply_personalization(self, model, personal_3dmm, listener_personality): - personal_3dmm = personal_3dmm.unsqueeze(0).to(self.device) if personal_3dmm.numel() > 0 else None - listener_personality = listener_personality.unsqueeze(0).to(self.device) - person_embedding = model.encode_person_condition( - p=personal_3dmm, - personality=listener_personality, - ) - model.kernel = model.hypernet(person_embedding) - model.apply_weights() - # (P4) at test the model is driven via model.main_net directly (not - # MainNetUnified.forward), so the coarse-plan FiLM must be stashed here too; - # it persists for every window of this person until the next person overwrites it. - if getattr(model, "person_coarse", None) is not None: - p_gamma, p_beta = model.person_coarse(person_embedding) - model._coarse_denoiser._person_coarse_film = (p_gamma, p_beta) - - @staticmethod - def _eeg_targets_from_motion_lengths(listener_eeg, listener_eeg_mask, motion_lengths): - if listener_eeg is None or listener_eeg.numel() == 0: - return None, None - if listener_eeg_mask is None or listener_eeg_mask.numel() == 0: - listener_eeg_mask = torch.ones_like(listener_eeg) - - indices = [] - offset = 0 - total_length = listener_eeg.shape[0] - for motion_length in motion_lengths: - length = int(motion_length.item() if torch.is_tensor(motion_length) else motion_length) - last_idx = min(max(offset + max(length, 1) - 1, 0), total_length - 1) - indices.append(last_idx) - offset += max(length, 0) - if not indices: - return None, None - - index_tensor = torch.tensor(indices, dtype=torch.long) - return ( - listener_eeg[index_tensor].unsqueeze(0).float(), - listener_eeg_mask[index_tensor].unsqueeze(0).float(), - ) - - def _predict_windows_ar(self, model, speaker_audio, speaker_emotion, speaker_3dmm, - motion_lengths): - """ONLINE autoregressive cross-window continuity (our generic-online recipe - ported here): generate windows in temporal order, feeding each window the - PREVIOUS window's own prediction as `past_listener_emotion`. With parallel - preds each prediction forms its own AR chain (past shape == num_preds, so the - matcher does NOT broadcast-replicate it). Window 0 gets no past. - """ - num_windows = speaker_audio.shape[0] - window_size = int(self.trainer_cfg.window_size) - predictions = [] - prev = None # (num_preds, window_size, d) previous-window prediction per chain - for wi in range(num_windows): - input_dict = { - "speaker_audio_input": speaker_audio[wi:wi + 1].to(self.device), - "speaker_emotion_input": speaker_emotion[wi:wi + 1].to(self.device), - "speaker_3dmm_input": speaker_3dmm[wi:wi + 1].to(self.device), - "motion_length": motion_lengths[wi:wi + 1].to(self.device), - "past_listener_emotion": prev, - } - outputs = model.main_net(**input_dict) - pred = outputs["prediction_emotion"].detach() # (1, num_preds, window_size, d) - predictions.append(pred.cpu()) - # carry this window's prediction as the next window's past (per chain). - nxt = pred[0] # (num_preds, window_size, d) - prev = nxt[:, -window_size:, :].to(self.device) - return torch.cat(predictions, dim=0) # (num_windows, num_preds, window_size, d) - - def _predict_windows_once(self, model, speaker_audio, speaker_emotion, speaker_3dmm, - motion_lengths, return_eeg=False): - predictions = [] - eeg_predictions = [] - total_windows = speaker_audio.shape[0] - for start in range(0, total_windows, self.eval_clip_batch_size): - end = start + self.eval_clip_batch_size - input_dict = { - "speaker_audio_input": speaker_audio[start:end].to(self.device), - "speaker_emotion_input": speaker_emotion[start:end].to(self.device), - "speaker_3dmm_input": speaker_3dmm[start:end].to(self.device), - "motion_length": motion_lengths[start:end].to(self.device), - } - outputs = model.main_net(**input_dict) - predictions.append(outputs["prediction_emotion"].detach().cpu()) - if return_eeg: - if "prediction_eeg" not in outputs: - raise RuntimeError( - "trainer.generic.eval_eeg=True but the model did not return prediction_eeg. " - "Check configs//model/motion_diffusion.yaml -> eeg_head.enabled." - ) - eeg_predictions.append(outputs["prediction_eeg"].detach().cpu()) - if return_eeg: - return torch.cat(predictions, dim=0), torch.cat(eeg_predictions, dim=0) - return torch.cat(predictions, dim=0) - - def _frrea_render_sample(self, renderer, latent_embedder, pred_listener_emotion, - listener_video, fake_dir, real_dir, stride, - shard_idx, batch_idx, sample_idx): - """(FRRea) Render one sample's generated (fake) + real listener frames for FID. - - Mirrors the generic trainer: renders prediction #0 onto the GT listener's - reference frame; real frames are the (already subsampled) GT listener clip. - Filenames are shard/batch/sample unique so parallel runs never collide.""" - import cv2 - lv = listener_video - if isinstance(lv, (list, tuple)): - lv = lv[0] - if not torch.is_tensor(lv) or lv.numel() == 0: - return - reference = lv[0].to(self.device) # (3, H, W) - emotion = pred_listener_emotion[0].to(self.device).float() # (clip_len, 25) - with torch.no_grad(): - coeff_3dmm = latent_embedder.decode_coeff(emotion) # (clip_len, 58) - fake_np, real_np = renderer.render_frames_for_fid( - coeff_3dmm, reference, lv, fake_stride=stride) - prefix = f"sh{shard_idx}_b{batch_idx}_s{sample_idx}" - for i in range(fake_np.shape[0]): - cv2.imwrite(os.path.join(fake_dir, f"{prefix}_f{i}.png"), fake_np[i]) - for i in range(real_np.shape[0]): - cv2.imwrite(os.path.join(real_dir, f"{prefix}_f{i}.png"), real_np[i]) - - def test(self): - logger.info("Loading rewrite-weight test data module") - self._ensure_eeg_data_enabled(stage="test") - test_loader = self.data_module.get_dataloader(stage="test") - logger.info("Test data module loaded") - - model = self._build_model(stage="test") - if not self.skip_modifier: - self._load_modifier_checkpoint(model, run_key="resume_runid", require_eeg_head=self.eval_eeg) - else: - model.main_net._identity_modifier = True - model.eval() - if model.person_encoder is not None: - model.person_encoder.eval() - - logger.info("Loading post processor") - post_processor = Processor( - config_name=self.post_config_name, - clip_len_test=self.post_clip_length, - device=self.device, - ) - logger.info("Post processor loaded") - - # ---- FRRea (FID) frame rendering setup (optional, gated by compute_frrea) ---- - compute_frrea = self._as_bool(self.kwargs.get("compute_frrea", False)) - frrea_renderer = frrea_latent_embedder = None - frrea_fake_dir = frrea_real_dir = None - frrea_stride = int(self.kwargs.get("frrea_stride", 30)) - frrea_shard_idx = int(os.environ.get("EVAL_SHARD_IDX", "0")) - if compute_frrea: - renderer_cfg = self.kwargs.get("renderer", None) - if renderer_cfg is None: - raise RuntimeError("compute_frrea=True requires trainer.renderer config.") - logger.info("Instantiating renderer for FRRea") - frrea_renderer = instantiate(renderer_cfg, device=self.device) - frrea_latent_embedder = model.main_net.diffusion_decoder.latent_embedder - tag = os.environ.get("FRREA_TAG", "rewrite") - base = os.path.join(hydra.utils.to_absolute_path("frrea_frames"), str(self.task), tag) - frrea_fake_dir = os.path.join(base, "fake") - frrea_real_dir = os.path.join(base, "real") - os.makedirs(frrea_fake_dir, exist_ok=True) - os.makedirs(frrea_real_dir, exist_ok=True) - logger.info(f"FRRea frames -> {base} (shard {frrea_shard_idx}, stride {frrea_stride})") - - gt_listener_emotions_all = [] - pred_listener_emotions_all = [] - input_speaker_emotions_all = [] - gt_listener_eeg_all = [] - pred_listener_eeg_all = [] - listener_eeg_mask_all = [] - - with torch.inference_mode(): - for frrea_batch_idx, batch in enumerate(tqdm(test_loader)): - listener_video_clips = None - if len(batch) == 13: - ( - speaker_audio_clips, - _, - speaker_emotion_clips, - speaker_3dmm_clips, - listener_video_clips, - listener_emotion_clips, - _listener_3dmm_clips, - personal_3dmm_clips, - listener_personality_clips, - listener_eeg_clips, - listener_eeg_mask_clips, - speaker_seq_lengths, - _listener_seq_lengths, - ) = batch - elif len(batch) == 11: - ( - speaker_audio_clips, - _, - speaker_emotion_clips, - speaker_3dmm_clips, - _, - listener_emotion_clips, - _listener_3dmm_clips, - personal_3dmm_clips, - listener_personality_clips, - speaker_seq_lengths, - _listener_seq_lengths, - ) = batch - listener_eeg_clips = listener_eeg_mask_clips = None - else: - ( - speaker_audio_clips, - _, - speaker_emotion_clips, - speaker_3dmm_clips, - _, - listener_emotion_clips, - _listener_3dmm_clips, - personal_3dmm_clips, - speaker_seq_lengths, - _listener_seq_lengths, - ) = batch - listener_personality_clips = [ - speaker_audio.new_zeros((0,)) - for speaker_audio in speaker_audio_clips - ] - listener_eeg_clips = listener_eeg_mask_clips = None - - if self.eval_eeg and listener_eeg_clips is None: - raise RuntimeError( - "trainer.generic.eval_eeg=True but the test dataloader did not return EEG labels." - ) - - eeg_clips = listener_eeg_clips if self.eval_eeg else [None] * len(speaker_audio_clips) - eeg_masks = listener_eeg_mask_clips if self.eval_eeg else [None] * len(speaker_audio_clips) - - for _frrea_si, (speaker_audio, speaker_emotion, speaker_3dmm, listener_gts, - personal_3dmm, listener_personality, listener_eeg, listener_eeg_mask, seq_length) in enumerate(zip( - speaker_audio_clips, - speaker_emotion_clips, - speaker_3dmm_clips, - listener_emotion_clips, - personal_3dmm_clips, - listener_personality_clips, - eeg_clips, - eeg_masks, - speaker_seq_lengths, - )): - length = int(seq_length.item() if torch.is_tensor(seq_length) else seq_length) - input_speaker_emotions_all.append(speaker_emotion[:length]) - gt_listener_emotions_all.append(listener_gts) - - windows_audio, windows_emotion, windows_3dmm, motion_lengths = self._build_test_windows( - speaker_audio, speaker_emotion, speaker_3dmm, length, - ) - if self.eval_eeg: - eeg_target, eeg_mask = self._eeg_targets_from_motion_lengths( - listener_eeg, listener_eeg_mask, motion_lengths, - ) - if eeg_target is None: - raise RuntimeError("EEG evaluation requested but a sample has no EEG target.") - - self._apply_personalization(model, personal_3dmm, listener_personality) - sample_predictions = [] - sample_eeg_predictions = [] - while len(sample_predictions) < self.num_eval_preds: - if self.eval_eeg: - window_predictions, window_eeg_predictions = self._predict_windows_once( - model, - windows_audio, - windows_emotion, - windows_3dmm, - motion_lengths, - return_eeg=True, - ) - elif self.task == "online" and self._as_bool( - self.trainer_cfg.get("online_autoregressive", True)): - window_predictions = self._predict_windows_ar( - model, - windows_audio, - windows_emotion, - windows_3dmm, - motion_lengths, - ) - window_eeg_predictions = None - else: - window_predictions = self._predict_windows_once( - model, - windows_audio, - windows_emotion, - windows_3dmm, - motion_lengths, - ) - window_eeg_predictions = None - sequence_predictions = rearrange( - window_predictions, - "b n w d -> n (b w) d", - )[:, :length] - sample_predictions.extend([prediction for prediction in sequence_predictions]) - if self.eval_eeg: - sequence_eeg_predictions = rearrange( - window_eeg_predictions, - "b n d -> n b d", - ) - sample_eeg_predictions.extend( - [prediction for prediction in sequence_eeg_predictions] - ) - - sample_prediction = torch.stack( - sample_predictions[:self.num_eval_preds], - dim=0, - ) - if self.data_clamp: - sample_prediction[:, :, :15] = torch.round(sample_prediction[:, :, :15]) - pred_listener_emotions_all.append(sample_prediction) - - if compute_frrea and listener_video_clips is not None: - self._frrea_render_sample( - frrea_renderer, frrea_latent_embedder, sample_prediction, - listener_video_clips[_frrea_si], frrea_fake_dir, frrea_real_dir, - frrea_stride, frrea_shard_idx, frrea_batch_idx, _frrea_si, - ) - - if self.eval_eeg: - pred_listener_eeg_all.append(torch.stack( - sample_eeg_predictions[:self.num_eval_preds], - dim=0, - )) - gt_listener_eeg_all.append(eeg_target) - listener_eeg_mask_all.append(eeg_mask) - - if len(pred_listener_emotions_all): - gt_listener_emotions_all = post_processor.forward( - prediction_list=pred_listener_emotions_all, - target_list=gt_listener_emotions_all, - ) - - results_to_save = { - "GT": gt_listener_emotions_all, - "PRED": pred_listener_emotions_all, - "INPUT": input_speaker_emotions_all, - } - if self.eval_eeg: - results_to_save.update({ - "GT_EEG": gt_listener_eeg_all, - "PRED_EEG": pred_listener_eeg_all, - "EEG_MASK": listener_eeg_mask_all, - }) - torch.save(results_to_save, "results.pt") - logger.info("Saved rewrite-weight results to results.pt") - - results = compute_metrics( - input_speaker_emotions_all, - pred_listener_emotions_all, - gt_listener_emotions_all, - ) - if self.eval_eeg: - results.update(compute_eeg_metrics( - pred_listener_eeg_all, - gt_listener_eeg_all, - listener_eeg_mask_all, - )) - logger.info(results) diff --git a/personalised/code/utils/logging.py b/personalised/code/utils/logging.py deleted file mode 100644 index 7c9e6f74e6ce9c7267e004474dbad0dfe91057c7..0000000000000000000000000000000000000000 --- a/personalised/code/utils/logging.py +++ /dev/null @@ -1,40 +0,0 @@ -import logging - - -class LevelsFilter(logging.Filter): - def __init__(self, levels): - self.levels = [getattr(logging, level) for level in levels] - - def filter(self, record): - return record.levelno in self.levels - - -class StreamToLogger(object): - """ - Fake file-like stream object that redirects writes to a logger instance. - """ - def __init__(self, logger, level): - self.logger = logger - self.level = level - self.linebuf = '' - - def write(self, buf): - for line in buf.rstrip().splitlines(): - self.logger.log(self.level, line.rstrip()) - - def flush(self): - pass - - -class TqdmLoggingHandler(logging.Handler): - def __init__(self, level=logging.NOTSET): - super().__init__(level) - - def emit(self, record): - import tqdm - try: - msg = self.format(record) - tqdm.tqdm.write(msg) - self.flush() - except Exception: - self.handleError(record) \ No newline at end of file diff --git a/personalised/code/utils/render.py b/personalised/code/utils/render.py deleted file mode 100644 index 81344fc3e6fb63cbf9cdb2fe4330119ed107fce1..0000000000000000000000000000000000000000 --- a/personalised/code/utils/render.py +++ /dev/null @@ -1,202 +0,0 @@ -import math -import os -import numpy as np -import torch -import hydra -from matplotlib import pyplot as plt -from torchvision import transforms -import cv2 -from utils.util import torch_img_to_np, _fix_image, torch_img_to_np2, set_seed -from external.FaceVerse import get_faceverse -from external.PIRender import FaceGenerator -from skimage.io import imsave -import skvideo.io - - -def obtain_seq_index(index, num_frames, semantic_radius=13): - seq = list(range(index - semantic_radius, index + semantic_radius + 1)) - seq = [min(max(item, 0), num_frames - 1) for item in seq] - return seq - - -def transform_semantic(semantic): - semantic_list = [] - for i in range(semantic.shape[0]): - index = obtain_seq_index(i, semantic.shape[0]) - semantic_item = semantic[index, :].unsqueeze(0) - semantic_list.append(semantic_item) - semantic = torch.cat(semantic_list, dim=0) - return semantic.transpose(1, 2) - - -class Render(object): - """Computes and stores the average and current value""" - - def __init__(self, device='cpu', **kwargs): - dir = hydra.utils.to_absolute_path("external/FaceVerse") - self.faceverse, _ = get_faceverse( - path=os.path.join(dir, "data/faceverse_simple_v2.npy"), - device=device, img_size=224) - self.faceverse.init_coeff_tensors() - self.id_tensor = torch.from_numpy(np.load(os.path.join(dir, "reference_full.npy"))).float().view(1, -1)[:, :150] - self.pi_render = FaceGenerator().to(device) - self.pi_render.eval() - ckpt_path = hydra.utils.to_absolute_path("external/PIRender/cur_model_fold.pth") - if not os.path.isfile(ckpt_path): - raise ValueError(f"No checkpoint found at {ckpt_path}") - pi_ckpt = torch.load(ckpt_path) - self.pi_render.load_state_dict(pi_ckpt['state_dict'] if 'state_dict' in pi_ckpt else pi_ckpt) - - self.mean_face = torch.FloatTensor( - np.load(os.path.join(dir, "mean_face.npy")).astype(np.float32)).view(1, 1, -1).to(device) - self.std_face = torch.FloatTensor( - np.load(os.path.join(dir, "std_face.npy")).astype(np.float32)).view(1, 1, -1).to(device) - - transform_reverse = kwargs.get('transform_reverse', 'zero_center') - if transform_reverse == 'zero_center': - self._reverse_transform_3dmm = transforms.Lambda(lambda e: e + self.mean_face) - elif transform_reverse == 'standard': - self._reverse_transform_3dmm = transforms.Lambda(lambda e: e * self.std_face + self.mean_face) - else: - raise ValueError(f"Unknown transform_reverse: {transform_reverse}") - self._transform = transforms.Lambda( - lambda e: (lambda tmp: tmp.__setitem__((slice(None), -1), e[:, -1] - self.mean_face[0, 0, -1]) or tmp)(e.clone())) - - def rendering(self, path, ind, listener_vectors, speaker_video_clip, listener_reference, listener_video_clip): - if len(listener_vectors.shape) > 2: - listener_vectors = listener_vectors.squeeze(0) - - # 3D video - T = listener_vectors.shape[0] - listener_vectors = self._reverse_transform_3dmm(listener_vectors)[0] - listener_vectors = self._transform(listener_vectors) - # print(f"maximum of listener_3dmm_out: {torch.max(listener_vectors)}") - # print(f"minimum of listener_3dmm_out: {torch.min(listener_vectors)}") - - T_unit = 512 - rendered_img_r_list = [] - for i in range(math.ceil(T / T_unit)): - if i != math.ceil(T / T_unit) - 1: - listener_vectors_i = listener_vectors[i * T_unit:(i + 1) * T_unit] - else: - listener_vectors_i = listener_vectors[i * T_unit:] - - self.faceverse.batch_size = listener_vectors_i.shape[0] - self.faceverse.init_coeff_tensors() - exp_tensor = listener_vectors_i[:, :52].to(listener_vectors_i.get_device()) - rot_tensor = listener_vectors_i[:, 52:55].to(listener_vectors_i.get_device()) - trans_tensor = listener_vectors_i[:, 55:].to(listener_vectors_i.get_device()) - self.faceverse.exp_tensor = exp_tensor - self.faceverse.rot_tensor = rot_tensor - self.faceverse.trans_tensor = trans_tensor - self.faceverse.id_tensor = self.id_tensor.reshape(1, 150).repeat( - listener_vectors_i.shape[0], 1).to(listener_vectors_i.get_device()) - - pred_dict = self.faceverse(self.faceverse.get_packed_tensors(), render=True, texture=False) - rendered_img_r = pred_dict['rendered_img'] - rendered_img_r = np.clip(rendered_img_r.cpu().numpy(), 0, 255) - rendered_img_r = rendered_img_r[:, :, :, :3].astype(np.uint8) - rendered_img_r_list.append(rendered_img_r) - rendered_img_r = np.concatenate(rendered_img_r_list, axis=0) - - # 2D video - semantics = transform_semantic(listener_vectors.detach()).to(listener_vectors.get_device()) - C, H, W = listener_reference.shape - output_dict_list = [] - duration = listener_vectors.shape[0] // 20 - listener_reference_frames = listener_reference.repeat(listener_vectors.shape[0], 1, 1).reshape( - listener_vectors.shape[0], C, H, W) - - for i in range(20): - if i != 19: - listener_reference_copy = listener_reference_frames[i * duration:(i + 1) * duration] - semantics_copy = semantics[i * duration:(i + 1) * duration] - else: - listener_reference_copy = listener_reference_frames[i * duration:] - semantics_copy = semantics[i * duration:] - with torch.no_grad(): - output_dict = self.pi_render(listener_reference_copy, semantics_copy) - fake_videos = output_dict['fake_image'] - fake_videos = torch_img_to_np2(fake_videos) - output_dict_list.append(fake_videos) - - listener_videos = np.concatenate(output_dict_list, axis=0) - speaker_video_clip = torch_img_to_np2(speaker_video_clip) - - out = cv2.VideoWriter(os.path.join(path, ind + "_val.avi"), cv2.VideoWriter_fourcc(*"MJPG"), 25, (672, 224)) - for i in range(rendered_img_r.shape[0]): - combined_img = np.zeros((224, 672, 3), dtype=np.uint8) - combined_img[0:224, 0:224] = speaker_video_clip[i] - combined_img[0:224, 224:448] = rendered_img_r[i] - combined_img[0:224, 448:] = listener_videos[i] - out.write(combined_img) - out.release() - - listener_video_clip = torch_img_to_np2(listener_video_clip) # [L, ...] - path_real = os.path.join(path, ind, 'real') - if not os.path.exists(path_real): - os.makedirs(path_real) - path_fake = os.path.join(path, ind, 'fake') - if not os.path.exists(path_fake): - os.makedirs(path_fake) - path_speaker = os.path.join(path, ind, 'speaker') - if not os.path.exists(path_speaker): - os.makedirs(path_speaker) - - # n_fake = listener_videos.shape[0] if hasattr(listener_videos, 'shape') else len(listener_videos) - # n_real = listener_video_clip.shape[0] if hasattr(listener_video_clip, 'shape') else len(listener_video_clip) - # n_speaker = speaker_video_clip.shape[0] if hasattr(speaker_video_clip, 'shape') else len(speaker_video_clip) - for i in range(0, rendered_img_r.shape[0], 30): - # if i < n_fake and i < n_real: - cv2.imwrite(os.path.join(path_fake, 'img_' + str(i + 1) + '.png'), listener_videos[i]) - cv2.imwrite(os.path.join(path_real, 'img_' + str(i + 1) + '.png'), listener_video_clip[i]) - cv2.imwrite(os.path.join(path_speaker, 'img_' + str(i + 1) + '.png'), speaker_video_clip[i]) - - def render_frames_for_fid(self, listener_vectors, listener_reference, - real_frames, fake_stride=30): - """Render generated (fake) listener frames via PIRender for the FRRea (FID) metric, - and return them alongside the (already-subsampled) real listener frames. - - The fake frames are produced from the predicted 3DMM sequence and subsampled by - ``fake_stride`` (FID needs a representative sample of frames, not every frame). - Temporal semantics are computed on the full predicted sequence first, so each - rendered frame still uses the correct +/-13 frame window. The real frames are - passed in already subsampled (decoded at stride by the dataset) and used as-is; - FID is distributional, so exact frame-by-frame alignment is not required. - - Args: - listener_vectors: predicted 3DMM coefficients (T, 58) in model output space. - listener_reference: a single real listener frame (3, H, W), normalized. - real_frames: already-subsampled real listener frames (M, 3, H, W), normalized. - fake_stride: keep one rendered frame every ``fake_stride`` predicted frames. - - Returns: - (fake_np, real_np): each (N, H, W, 3) uint8 BGR arrays. - """ - device = listener_reference.device - if len(listener_vectors.shape) > 2: - listener_vectors = listener_vectors.squeeze(0) - listener_vectors = listener_vectors.to(device) - T = listener_vectors.shape[0] - - listener_vectors = self._reverse_transform_3dmm(listener_vectors)[0] - listener_vectors = self._transform(listener_vectors) - - # Full-sequence semantics so per-frame +/-13 windows are correct, then subsample. - semantics = transform_semantic(listener_vectors.detach()).to(device) - - idx = list(range(0, T, fake_stride)) - C, H, W = listener_reference.shape - ref = listener_reference.unsqueeze(0).repeat(len(idx), 1, 1, 1) - sem_sel = semantics[idx] - - fake_list = [] - chunk = 64 - with torch.no_grad(): - for i in range(0, len(idx), chunk): - out = self.pi_render(ref[i:i + chunk], sem_sel[i:i + chunk]) - fake_list.append(torch_img_to_np2(out['fake_image'])) - fake_np = np.concatenate(fake_list, axis=0) - - real_np = torch_img_to_np2(real_frames.to(device)) - return fake_np, real_np \ No newline at end of file diff --git a/personalised/code/utils/runid.py b/personalised/code/utils/runid.py deleted file mode 100644 index cbfb68b92907684e6834e2677cd08c2b297aeac2..0000000000000000000000000000000000000000 --- a/personalised/code/utils/runid.py +++ /dev/null @@ -1,18 +0,0 @@ -""" -runid util. -Taken from wandb.sdk.lib.runid -""" - -from datetime import datetime -import shortuuid # type: ignore - - -def generate_id() -> str: - # ~3t run ids (36**8) - run_gen = shortuuid.ShortUUID(alphabet=list("0123456789abcdefghijklmnopqrstuvwxyz")) - - # prepend time string to 'run_gen' - prefix = datetime.now().strftime("%y%m%d%H%M%S") - - return prefix + '_' + run_gen.random(8) - # return run_gen.random(8) diff --git a/personalised/code/utils/util.py b/personalised/code/utils/util.py deleted file mode 100644 index ef1b155b477ad768357189419a69102c19f2ff97..0000000000000000000000000000000000000000 --- a/personalised/code/utils/util.py +++ /dev/null @@ -1,220 +0,0 @@ -from datetime import datetime -from math import cos, pi -from torchvision import transforms -from PIL import Image -import torch.nn as nn -import cv2 -import torch.nn.functional as F -from omegaconf import OmegaConf -import os -import yaml -from torch.backends import cudnn -import random -import numpy as np -import torch -import json -from framework.metrics.metric import baseline_random, baseline_mime -from framework.utils.compute_metrics import compute_metrics - - -def set_seed(seed: int, deterministic: bool = True) -> None: - # Python - random.seed(seed) - # NumPy - np.random.seed(seed) - # PyTorch CPU - torch.manual_seed(seed) - # PyTorch GPU - if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - if deterministic: - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False - - -def load_config(config_path=None): - cli_conf = OmegaConf.from_cli() - model_conf = OmegaConf.load(cli_conf.pop('config') if config_path is None else config_path) - return OmegaConf.merge(model_conf, cli_conf) - - -def load_config_from_file(path): - return OmegaConf.load(path) - - -def get_logging_path(log_dir): - current_time = datetime.now() - time_str = str(current_time) - time_str = '-'.join(time_str.split(' ')) - time_str = time_str.split('.')[0] - lod_dir = os.path.join(log_dir, time_str) - return lod_dir - - -def get_tensorboard_path(tb_dir): - current_time = datetime.now() - time_str = str(current_time) - time_str = '-'.join(time_str.split(' ')) - time_str = time_str.split('.')[0] - tb_dir = os.path.join(tb_dir, time_str) - os.makedirs(tb_dir, exist_ok=True) - return tb_dir - - -def store_config(config): - # store config to directory - dir = config.trainer.out_dir - os.makedirs(dir, exist_ok=True) - with open(os.path.join(dir, "config.yaml"), "w") as f: - yaml.dump(OmegaConf.to_container(config), f) - - -def torch_img_to_np(img): - return img.detach().cpu().numpy().transpose(0, 2, 3, 1) - - -def torch_img_to_np2(img): - img = img.detach().cpu().numpy() - # img = img * np.array([0.229, 0.224, 0.225]).reshape(1,-1,1,1) - # img = img + np.array([0.485, 0.456, 0.406]).reshape(1,-1,1,1) - img = img * np.array([0.5, 0.5, 0.5]).reshape(1, -1, 1, 1) - img = img + np.array([0.5, 0.5, 0.5]).reshape(1, -1, 1, 1) - img = img.transpose(0, 2, 3, 1) - img = img * 255.0 - img = np.clip(img, 0, 255).astype(np.uint8)[:, :, :, [2, 1, 0]] - - return img - - -def _fix_image(image): - if image.max() < 30.: - image = image * 255. - image = np.clip(image, 0, 255).astype(np.uint8)[:, :, :, [2, 1, 0]] - return image - - -class AverageMeter(object): - """Computes and stores the average and current value""" - - def __init__(self): - self.reset() - - def reset(self): - self.val = 0 - self.avg = 0 - self.sum = 0 - self.count = 0 - - def update(self, val, n=1): - self.val = val - self.sum += val * n - self.count += n - self.avg = self.sum / self.count - - -def accuracy(output, target, topk=(1,)): - """Computes the precision@k for the specified values of k""" - maxk = max(topk) - batch_size = target.size(0) - - _, pred = output.topk(maxk, 1, True, True) - pred = pred.t() - correct = pred.eq(target.view(1, -1).expand_as(pred)) - - res = [] - for k in topk: - correct_k = correct[:k].view(-1).float().sum(0) - res.append(correct_k.mul_(100.0 / batch_size)) - return res - - -def binary_accuracy(pred, labels): - """Computes the precision@k for the binary classification""" - pred = (pred.cpu().detach().numpy() > 0.5).astype(int) - labels = labels.cpu().detach().numpy() - - correct = np.sum(pred == labels) - return correct / len(labels) - - -def get_lr(optimizer): - for param_group in optimizer.param_groups: - return param_group['lr'] - - -def save_results(results, filename='results.json'): - with open(filename, 'w') as f: - json.dump(results, f, indent=4) - - -if __name__ == '__main__': - speaker_inputs = torch.load("/home/x/xk18/react-challange/main/react-challenge-2025/" - "results/speaker_inputs.pt")["speaker"] - print(f"Length of speaker_inputs: {len(speaker_inputs)}") - print(f"shape of speaker_inputs[0]: {speaker_inputs[0].shape}") - diffusion_online_result = torch.load("/home/x/xk18/react-challange/main/react-challenge-2025/results/" - "motion_diffusion_online_results.pt") - listener_targets = diffusion_online_result["GT"] - print(f"Length of listener_targets: {len(listener_targets)}") - print(f"shape of listener_targets[0]: {listener_targets[0].shape}") - - # 1. random prediction - print("Random prediction:") - try: - listener_predictions = [baseline_random(listener_target) for listener_target in listener_targets] - results = compute_metrics(speaker_inputs=speaker_inputs, - listener_predictions=listener_predictions, - listener_targets=listener_targets) - save_results(results, filename="random_results.json") - except Exception as e: - print(e) - print("Random prediction failed") - - # 2. mimic speaker - print("Mimic speaker:") - try: - listener_predictions = [baseline_mime(speaker_input) for speaker_input in speaker_inputs] - results = compute_metrics(speaker_inputs=speaker_inputs, - listener_predictions=listener_predictions, - listener_targets=listener_targets) - save_results(results, filename="mime_results.json") - except Exception as e: - print(e) - print("Mimic speaker failed") - - # 3. GT identical - print("GT identical:") - try: - listener_predictions = listener_targets - results = compute_metrics(speaker_inputs=speaker_inputs, - listener_predictions=listener_predictions, - listener_targets=listener_targets) - save_results(results, filename="GT_results.json") - except Exception as e: - print(e) - print("GT identical failed") - - # 4. Mean Frame - print("Mean Frame:") - try: - mean_frame = torch.tensor([0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, -0.07756615, 0.14737537, - 0.2976104, 0.25274006, 0.1022017, 0.08290554, 0.0072741, 0.15682219, 0.03406566, - 0.06638012]) - mean_frame_expanded = mean_frame[None, None, :] - - listener_predictions = [] - for target in listener_targets: - pred = torch.tile( - mean_frame_expanded, - (target.shape[0], target.shape[1], 1) - ) - listener_predictions.append(pred) - - results = compute_metrics(speaker_inputs=speaker_inputs, - listener_predictions=listener_predictions, - listener_targets=listener_targets) - save_results(results, filename="mean_frame_results.json") - except Exception as e: - print(e) - print("Mean Frame failed") \ No newline at end of file