diff --git a/personalised/code/configs/README.md b/personalised/code/configs/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cf16b6458d7461b93c15e44353b1cde2fad05975 --- /dev/null +++ b/personalised/code/configs/README.md @@ -0,0 +1,27 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..c1578cc203b43cf4f693650cea47a429b13937df --- /dev/null +++ b/personalised/code/configs/data/emotion_autoencoder.yaml @@ -0,0 +1,29 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..05d60f88d8974e59c20120be88a05c0d84d4fc4e --- /dev/null +++ b/personalised/code/configs/data/motion_diffusion.yaml @@ -0,0 +1,122 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..b6b2c154a346926125a29a335450fb6b2a5193c6 --- /dev/null +++ b/personalised/code/configs/data/motion_transvae.yaml @@ -0,0 +1,125 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..06e8e5fe9131535e98886a7e3f694dc46023a1c8 --- /dev/null +++ b/personalised/code/configs/data/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,135 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..220fef845c2016a301e462bc5639386abf7fa0ea --- /dev/null +++ b/personalised/code/configs/g2p_delta.yaml @@ -0,0 +1,88 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..05d60f88d8974e59c20120be88a05c0d84d4fc4e --- /dev/null +++ b/personalised/code/configs/generic_offline/data/motion_diffusion.yaml @@ -0,0 +1,122 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..b6b2c154a346926125a29a335450fb6b2a5193c6 --- /dev/null +++ b/personalised/code/configs/generic_offline/data/motion_transvae.yaml @@ -0,0 +1,125 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..f528b1e1cb4be34a38fde02543cc3945f2dd520f --- /dev/null +++ b/personalised/code/configs/generic_offline/model/losses/motion_diffusion.yaml @@ -0,0 +1,11 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..fd96f6674d5a4b7ddfc2caf12a95c19b4ef4d7ea --- /dev/null +++ b/personalised/code/configs/generic_offline/model/losses/motion_transvae.yaml @@ -0,0 +1,8 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..59385f65e8a2d603bc2c32b5abea0e8ee481cc02 --- /dev/null +++ b/personalised/code/configs/generic_offline/model/motion_diffusion.yaml @@ -0,0 +1,116 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..8e18655a404101b4ee85faa3956a14c00ea1cc4b --- /dev/null +++ b/personalised/code/configs/generic_offline/model/motion_diffusion/audio_embedder.yaml @@ -0,0 +1,4 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..08e0afc25b4cadc9af577a8faefb0761cdd8f722 --- /dev/null +++ b/personalised/code/configs/generic_offline/model/motion_transvae.yaml @@ -0,0 +1,24 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..61651ade05e88d28af3ffa4e65d7618fc49fde41 --- /dev/null +++ b/personalised/code/configs/generic_offline/motion_diffusion.yaml @@ -0,0 +1,23 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..47f14222e2a4fd1fad4ec94afb22b927cb900cf8 --- /dev/null +++ b/personalised/code/configs/generic_offline/motion_transvae.yaml @@ -0,0 +1,23 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..b2ca1d565e81819eaeac2485c2b2bf3f7c10bf87 --- /dev/null +++ b/personalised/code/configs/generic_offline/trainer/motion_diffusion.yaml @@ -0,0 +1,48 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..2a7768d1c16df8375a0075a0619734cf0fe4ca03 --- /dev/null +++ b/personalised/code/configs/generic_offline/trainer/motion_transvae.yaml @@ -0,0 +1,42 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..05d60f88d8974e59c20120be88a05c0d84d4fc4e --- /dev/null +++ b/personalised/code/configs/generic_online/data/motion_diffusion.yaml @@ -0,0 +1,122 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..b6b2c154a346926125a29a335450fb6b2a5193c6 --- /dev/null +++ b/personalised/code/configs/generic_online/data/motion_transvae.yaml @@ -0,0 +1,125 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..f528b1e1cb4be34a38fde02543cc3945f2dd520f --- /dev/null +++ b/personalised/code/configs/generic_online/model/losses/motion_diffusion.yaml @@ -0,0 +1,11 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..fd96f6674d5a4b7ddfc2caf12a95c19b4ef4d7ea --- /dev/null +++ b/personalised/code/configs/generic_online/model/losses/motion_transvae.yaml @@ -0,0 +1,8 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..eb29cb006f990b7bcaa3a2a238aa8e0633771967 --- /dev/null +++ b/personalised/code/configs/generic_online/model/motion_diffusion.yaml @@ -0,0 +1,116 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..8e18655a404101b4ee85faa3956a14c00ea1cc4b --- /dev/null +++ b/personalised/code/configs/generic_online/model/motion_diffusion/audio_embedder.yaml @@ -0,0 +1,4 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..08e0afc25b4cadc9af577a8faefb0761cdd8f722 --- /dev/null +++ b/personalised/code/configs/generic_online/model/motion_transvae.yaml @@ -0,0 +1,24 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..0d2d17d2d32413cb9c153824b5fee5ccdd074dd9 --- /dev/null +++ b/personalised/code/configs/generic_online/motion_diffusion.yaml @@ -0,0 +1,23 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..dc6398335044355b2311881bff216d4d42983784 --- /dev/null +++ b/personalised/code/configs/generic_online/motion_transvae.yaml @@ -0,0 +1,23 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..0712883a80805af1dde081226de910e6b64abb9f --- /dev/null +++ b/personalised/code/configs/generic_online/trainer/motion_diffusion.yaml @@ -0,0 +1,48 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..bcff8934156353f9474d99413ae2c81e23fbe640 --- /dev/null +++ b/personalised/code/configs/generic_online/trainer/motion_transvae.yaml @@ -0,0 +1,42 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..3da2aede1c03e6d9660ccabe8ce187288aee3f48 --- /dev/null +++ b/personalised/code/configs/logger/none.yaml @@ -0,0 +1,4 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..36e25fe9263b0dfb05a98eb37070e36e6f13e9db --- /dev/null +++ b/personalised/code/configs/main.yaml @@ -0,0 +1,26 @@ + +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 new file mode 100644 index 0000000000000000000000000000000000000000..2f80bbb7189d799d036a01b075ba3cc332a83b03 --- /dev/null +++ b/personalised/code/configs/model/emotion_autoencoder.yaml @@ -0,0 +1,17 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..6881880215136f33e15ae18141d3470b69d6afc2 --- /dev/null +++ b/personalised/code/configs/model/losses/emotion_autoencoder.yaml @@ -0,0 +1,6 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..f528b1e1cb4be34a38fde02543cc3945f2dd520f --- /dev/null +++ b/personalised/code/configs/model/losses/motion_diffusion.yaml @@ -0,0 +1,11 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..9bd1b7aca31f445331118ddb566a2222d1fd3f31 --- /dev/null +++ b/personalised/code/configs/model/losses/motion_diffusion_causal.yaml @@ -0,0 +1,9 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..fd96f6674d5a4b7ddfc2caf12a95c19b4ef4d7ea --- /dev/null +++ b/personalised/code/configs/model/losses/motion_transvae.yaml @@ -0,0 +1,8 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..a8cd2da396773e073c381619488d0031e429476a --- /dev/null +++ b/personalised/code/configs/model/losses/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,8 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..fd313fcd0a1c96e3bb90552b6a3b51b14c336e05 --- /dev/null +++ b/personalised/code/configs/model/motion_diffusion.yaml @@ -0,0 +1,116 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..8e18655a404101b4ee85faa3956a14c00ea1cc4b --- /dev/null +++ b/personalised/code/configs/model/motion_diffusion/audio_embedder.yaml @@ -0,0 +1,4 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..1317b290021f3a19e682a75d4fa53142b942a3f6 --- /dev/null +++ b/personalised/code/configs/model/motion_diffusion_causal.yaml @@ -0,0 +1,20 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..d207d42dacd170cb509dcac0c4d5876edf9c6eda --- /dev/null +++ b/personalised/code/configs/model/motion_diffusion_velocity.yaml @@ -0,0 +1,116 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..08e0afc25b4cadc9af577a8faefb0761cdd8f722 --- /dev/null +++ b/personalised/code/configs/model/motion_transvae.yaml @@ -0,0 +1,24 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..f751f0d1d7f36e371fe68e3f45fdd5c6991c0b9d --- /dev/null +++ b/personalised/code/configs/model/optim/adam.yaml @@ -0,0 +1,5 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..6c06aa42da9169cc696942320b80ead3e73206a7 --- /dev/null +++ b/personalised/code/configs/model/optim/adamw.yaml @@ -0,0 +1,7 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..18831c513bded4af228da868f302e6f12e165559 --- /dev/null +++ b/personalised/code/configs/model/optim/sgd.yaml @@ -0,0 +1,5 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..4178cbbf8eb8d4ed7ba95d463d18f49075e6db79 --- /dev/null +++ b/personalised/code/configs/model/scheduler/cosine_annealing.yaml @@ -0,0 +1,4 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..e510054eda62eee163e6383cd716c54b226ce784 --- /dev/null +++ b/personalised/code/configs/path.yaml @@ -0,0 +1,5 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..5ed98c9e0b85c0ef38972c73e741cceb0a603fc1 --- /dev/null +++ b/personalised/code/configs/pers_offline_causal.yaml @@ -0,0 +1,20 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..2f2e3e37417144295566cd16932153b5ca605bcb --- /dev/null +++ b/personalised/code/configs/pers_online_causal.yaml @@ -0,0 +1,19 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..06e8e5fe9131535e98886a7e3f694dc46023a1c8 --- /dev/null +++ b/personalised/code/configs/personalized_offline/data/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,135 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..f528b1e1cb4be34a38fde02543cc3945f2dd520f --- /dev/null +++ b/personalised/code/configs/personalized_offline/model/losses/motion_diffusion.yaml @@ -0,0 +1,11 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..a8cd2da396773e073c381619488d0031e429476a --- /dev/null +++ b/personalised/code/configs/personalized_offline/model/losses/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,8 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..7dcde62f171fefd7caa03974090625e9af068ff2 --- /dev/null +++ b/personalised/code/configs/personalized_offline/model/motion_diffusion.yaml @@ -0,0 +1,116 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..8e18655a404101b4ee85faa3956a14c00ea1cc4b --- /dev/null +++ b/personalised/code/configs/personalized_offline/model/motion_diffusion/audio_embedder.yaml @@ -0,0 +1,4 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..d4c9ab4249c5046bfcd830950389218e3d278aac --- /dev/null +++ b/personalised/code/configs/personalized_offline/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,23 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..a7d28d15d69518a5fa44b17e0ae4695c87817bf8 --- /dev/null +++ b/personalised/code/configs/personalized_offline/trainer/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,89 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..06e8e5fe9131535e98886a7e3f694dc46023a1c8 --- /dev/null +++ b/personalised/code/configs/personalized_online/data/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,135 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..f528b1e1cb4be34a38fde02543cc3945f2dd520f --- /dev/null +++ b/personalised/code/configs/personalized_online/model/losses/motion_diffusion.yaml @@ -0,0 +1,11 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..a8cd2da396773e073c381619488d0031e429476a --- /dev/null +++ b/personalised/code/configs/personalized_online/model/losses/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,8 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..3bddb28da79a5570ca657a268db3e9b4f484b311 --- /dev/null +++ b/personalised/code/configs/personalized_online/model/motion_diffusion.yaml @@ -0,0 +1,116 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..8e18655a404101b4ee85faa3956a14c00ea1cc4b --- /dev/null +++ b/personalised/code/configs/personalized_online/model/motion_diffusion/audio_embedder.yaml @@ -0,0 +1,4 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..c50bacd98086a6d0354ba1b620da7f7805de182a --- /dev/null +++ b/personalised/code/configs/personalized_online/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,23 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..0129543e100797b74c0cf6c2616f772c527c78bc --- /dev/null +++ b/personalised/code/configs/personalized_online/trainer/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,89 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..c1578cc203b43cf4f693650cea47a429b13937df --- /dev/null +++ b/personalised/code/configs/shared/data/emotion_autoencoder.yaml @@ -0,0 +1,29 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..3da2aede1c03e6d9660ccabe8ce187288aee3f48 --- /dev/null +++ b/personalised/code/configs/shared/logger/none.yaml @@ -0,0 +1,4 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..2f80bbb7189d799d036a01b075ba3cc332a83b03 --- /dev/null +++ b/personalised/code/configs/shared/model/emotion_autoencoder.yaml @@ -0,0 +1,17 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..6881880215136f33e15ae18141d3470b69d6afc2 --- /dev/null +++ b/personalised/code/configs/shared/model/losses/emotion_autoencoder.yaml @@ -0,0 +1,6 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..f751f0d1d7f36e371fe68e3f45fdd5c6991c0b9d --- /dev/null +++ b/personalised/code/configs/shared/model/optim/adam.yaml @@ -0,0 +1,5 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..6c06aa42da9169cc696942320b80ead3e73206a7 --- /dev/null +++ b/personalised/code/configs/shared/model/optim/adamw.yaml @@ -0,0 +1,7 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..18831c513bded4af228da868f302e6f12e165559 --- /dev/null +++ b/personalised/code/configs/shared/model/optim/sgd.yaml @@ -0,0 +1,5 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..4178cbbf8eb8d4ed7ba95d463d18f49075e6db79 --- /dev/null +++ b/personalised/code/configs/shared/model/scheduler/cosine_annealing.yaml @@ -0,0 +1,4 @@ +_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 new file mode 100644 index 0000000000000000000000000000000000000000..e510054eda62eee163e6383cd716c54b226ce784 --- /dev/null +++ b/personalised/code/configs/shared/path.yaml @@ -0,0 +1,5 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..33721984a93040d270b2a96608b75ea1db4ed9b4 --- /dev/null +++ b/personalised/code/configs/shared/trainer/emotion_autoencoder.yaml @@ -0,0 +1,23 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..703c7e0c9966ae65665657a43c33cf98360e8ee1 --- /dev/null +++ b/personalised/code/configs/trainer/_archive.yaml @@ -0,0 +1,105 @@ +## 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 new file mode 100644 index 0000000000000000000000000000000000000000..4df0d9bfb6ea69c75f8b0830827743dac16c0759 --- /dev/null +++ b/personalised/code/configs/trainer/emotion_autoencoder.yaml @@ -0,0 +1,23 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..b5ce777f5f7b9ecedbfc484ad7de0be352f25c37 --- /dev/null +++ b/personalised/code/configs/trainer/motion_diffusion.yaml @@ -0,0 +1,48 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..8380d6000738d4a50155cf6712a4ec7a0093017a --- /dev/null +++ b/personalised/code/configs/trainer/motion_diffusion_causal.yaml @@ -0,0 +1,48 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..7acbdc23912ffcd6421e6dd76f45b5b9a60abc53 --- /dev/null +++ b/personalised/code/configs/trainer/motion_diffusion_causal_dp.yaml @@ -0,0 +1,48 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..30348d360434e79b4f261ee4a262448916b17b2b --- /dev/null +++ b/personalised/code/configs/trainer/motion_diffusion_dp.yaml @@ -0,0 +1,48 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..7d2a62f40ec2b46642c093b78ce02b99b20a1f67 --- /dev/null +++ b/personalised/code/configs/trainer/motion_diffusion_velocity.yaml @@ -0,0 +1,48 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..3069fac3d0e465f1733827a8953abdb270a425bb --- /dev/null +++ b/personalised/code/configs/trainer/motion_transvae.yaml @@ -0,0 +1,42 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..317d04df69f2b14725b5c5d91ec528d88cd34509 --- /dev/null +++ b/personalised/code/configs/trainer/perfrdiff_rewrite_weight.yaml @@ -0,0 +1,89 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/personalised/code/dataset/data_preprocess/audio_feature_extraction.py b/personalised/code/dataset/data_preprocess/audio_feature_extraction.py new file mode 100644 index 0000000000000000000000000000000000000000..bb3f7c94f0a36399dc6f81ce3019d35f84248862 --- /dev/null +++ b/personalised/code/dataset/data_preprocess/audio_feature_extraction.py @@ -0,0 +1,76 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..7827c031ec70e08de3e193f261f370194efdc540 --- /dev/null +++ b/personalised/code/dataset/data_preprocess/audio_feature_extraction_test.py @@ -0,0 +1,120 @@ +""" +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 new file mode 100644 index 0000000000000000000000000000000000000000..792207f3bf7bce9950528383b339081af15f49ca --- /dev/null +++ b/personalised/code/dataset/modules/audio_embedder.py @@ -0,0 +1,16 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..faa1117ea062b6ef9510ead08e3079905fc56006 --- /dev/null +++ b/personalised/code/dataset/modules/audio_processor.py @@ -0,0 +1,189 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..7a9062b2531fd86d9ef120670397512d3c593739 --- /dev/null +++ b/personalised/code/dataset/perfrdiff_rewrite_weight.py @@ -0,0 +1,534 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..713818969f3ba13bec3c688f115939266f9232f5 --- /dev/null +++ b/personalised/code/dataset/react_2025.py @@ -0,0 +1,621 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..7594b86e1205110152f6fbbf3dc0a234d92e5b17 --- /dev/null +++ b/personalised/code/dataset/tools/util.py @@ -0,0 +1,70 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..5345641eeccf3966467f44896f87f15882eba173 --- /dev/null +++ b/personalised/code/external/FaceVerse/FaceVerseModel.py @@ -0,0 +1,321 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..ba7be5c07c029703156ea22a68e9d1a62dc796bc --- /dev/null +++ b/personalised/code/external/FaceVerse/LICENSE @@ -0,0 +1,21 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..c56aa8337fd9a100a2d6c47857e61e4cf7d7d98e --- /dev/null +++ b/personalised/code/external/FaceVerse/ModelRenderer.py @@ -0,0 +1,60 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..38ba6ca7563a27fc9dc1948b04c4248c9cbfee88 --- /dev/null +++ b/personalised/code/external/FaceVerse/__init__.py @@ -0,0 +1,13 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..53726c387cad7910c858bda21304090d4c6d9899 --- /dev/null +++ b/personalised/code/external/FaceVerse/data/faceverse_simple_v2.npy @@ -0,0 +1,3 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..5140b79c85a712a2cb60d8eb809659d5f03d2cb8 --- /dev/null +++ b/personalised/code/external/FaceVerse/mean_face.npy @@ -0,0 +1,3 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..1fdd7e9e77db464b019173bbc1cfeb2f36a5f0b3 --- /dev/null +++ b/personalised/code/external/FaceVerse/reference_full.npy @@ -0,0 +1,3 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..8614b369ebba91b0c87c4a54d3ee433359d8f886 --- /dev/null +++ b/personalised/code/external/FaceVerse/std_face.npy @@ -0,0 +1,3 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..14ff991e09e38a074a9eb730f19d88e7a9f2c3b7 --- /dev/null +++ b/personalised/code/framework/feature_extractor/wav2vec.py @@ -0,0 +1,216 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..404ec802c7be3664b1d4a65fdedf54460627663c --- /dev/null +++ b/personalised/code/framework/g2p_delta/__init__.py @@ -0,0 +1,3 @@ +from .model import G2PDeltaModel + +__all__ = ["G2PDeltaModel"] diff --git a/personalised/code/framework/g2p_delta/model.py b/personalised/code/framework/g2p_delta/model.py new file mode 100644 index 0000000000000000000000000000000000000000..aac7e555b55569b2709cb67183be849960dec138 --- /dev/null +++ b/personalised/code/framework/g2p_delta/model.py @@ -0,0 +1,325 @@ +"""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 new file mode 100644 index 0000000000000000000000000000000000000000..54d72e19d21c9da1848f121d8f3c53701281be79 --- /dev/null +++ b/personalised/code/framework/metrics/ACC.py @@ -0,0 +1,4 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..e296198cb82b8b071f02dea918aeefeb8e59e8ea --- /dev/null +++ b/personalised/code/framework/metrics/FID.py @@ -0,0 +1,140 @@ +"""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 new file mode 100644 index 0000000000000000000000000000000000000000..716d876f0828b0e1b031cb409e330eeb3a7167d9 --- /dev/null +++ b/personalised/code/framework/metrics/FRC.py @@ -0,0 +1,208 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..c5ee673fa2f1d5bcd4b6ad116677fae720193529 --- /dev/null +++ b/personalised/code/framework/metrics/FRD.py @@ -0,0 +1,45 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..5285bba82ba24ae38734602019a3372cfa6dd26f --- /dev/null +++ b/personalised/code/framework/metrics/FRDvs.py @@ -0,0 +1,33 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..2333548fb622937985a7010b7c8b4e1739072739 --- /dev/null +++ b/personalised/code/framework/metrics/FRVar.py @@ -0,0 +1,24 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..5cdbbdfaaf731d0fd775fc958df2431a79b63e18 --- /dev/null +++ b/personalised/code/framework/metrics/S_MSE.py @@ -0,0 +1,25 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..377868a309c7f279df8cf0bbdaa58cd22d9fcbff --- /dev/null +++ b/personalised/code/framework/metrics/TLCC.py @@ -0,0 +1,101 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..94bb02c15266d4dc5f328140d8ccdf037426c991 --- /dev/null +++ b/personalised/code/framework/metrics/__init__.py @@ -0,0 +1,6 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..9cb8995a5139a7941f0065ff4d71c8062123280c --- /dev/null +++ b/personalised/code/framework/metrics/metric.py @@ -0,0 +1,90 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/personalised/code/framework/modules/emotion_autoencoder.py b/personalised/code/framework/modules/emotion_autoencoder.py new file mode 100644 index 0000000000000000000000000000000000000000..b4e0fe1c062762924106999e245a2e380cdf2277 --- /dev/null +++ b/personalised/code/framework/modules/emotion_autoencoder.py @@ -0,0 +1,310 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..c00c62a4dc6b6cc01f1d1fbe4715fc5624d4cbf5 --- /dev/null +++ b/personalised/code/framework/modules/post_processor.py @@ -0,0 +1,149 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/__init__.py b/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 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 new file mode 100644 index 0000000000000000000000000000000000000000..33f25255110cccb9927947118c803530ab64db2b --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/transformer_denoiser.py @@ -0,0 +1,562 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..8747a0c7e8d933f43e44475504753f8ac992e3bf --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/diffusion_decoder/transformer_denoiser_causal.py @@ -0,0 +1,241 @@ +""" +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 new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 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 new file mode 100644 index 0000000000000000000000000000000000000000..13746c145860e110d562164a5e2ee315f023a738 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/diffusion_prior/rotary_embedding_torch.py @@ -0,0 +1,282 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..7cee0e0159288a52609755f203fe88bacc3c553b --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/diffusion_prior/transformer_prior.py @@ -0,0 +1,435 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..61c8276570aba103b10403e505135d1cb3789cb4 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/gaussian_diffusion.py @@ -0,0 +1,1161 @@ +""" +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 new file mode 100644 index 0000000000000000000000000000000000000000..c14cf2ddf21d53070e4ce0bc689a282c97d6dc58 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/matchers.py @@ -0,0 +1,719 @@ +""" +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 new file mode 100644 index 0000000000000000000000000000000000000000..8219b5c990324e9d256e904a8ec8b2e51015b473 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/matchers_causal.py @@ -0,0 +1,125 @@ +""" +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 new file mode 100644 index 0000000000000000000000000000000000000000..f0bd206541468bcd048e1c68de305d449f8e8675 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/matchers_velocity.py @@ -0,0 +1,122 @@ +"""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 new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/personalised/code/framework/motion_diffusion/diffusion/operator/cross_attention.py b/personalised/code/framework/motion_diffusion/diffusion/operator/cross_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..e7b6e777f1945451e0c65b8b949c18908a00192f --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/operator/cross_attention.py @@ -0,0 +1,416 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..179582bd988bc996c2d056f739af298c1da9fd06 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/operator/embeddings.py @@ -0,0 +1,320 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..a749559173c88ba4401ba736890393fb42fa538a --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/operator/position_encoding.py @@ -0,0 +1,179 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..c067cfdb09326c7fc2b399da707695c556af3334 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/resample.py @@ -0,0 +1,153 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..53532c2141f9d3288b9077acc2b8a91c890954b9 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/rnn.py @@ -0,0 +1,282 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..d1bcd5c5fc0570d5eadfbeb10473b262f33b6f03 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/torch.py @@ -0,0 +1,242 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/personalised/code/framework/motion_diffusion/diffusion/utils/geometry.py b/personalised/code/framework/motion_diffusion/diffusion/utils/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..71b1c3caa86d6bf570b6b8f1edaa3549bade7bfc --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/utils/geometry.py @@ -0,0 +1,473 @@ +# -*- 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 new file mode 100644 index 0000000000000000000000000000000000000000..3e757bc5c8189e52b574e9fd859cdc00cf0503ea --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/utils/losses.py @@ -0,0 +1,69 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..717ed8f3ddbbecda431a5495d9f410db678899c9 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/utils/temos_utils.py @@ -0,0 +1,131 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..5fac735a73ef9958af99718ad6d8c7389d74cebe --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/utils/util.py @@ -0,0 +1,69 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..ac422df10f4965a9d2971350b98c01393199b1d5 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/diffusion/velocity_transform.py @@ -0,0 +1,41 @@ +"""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 new file mode 100644 index 0000000000000000000000000000000000000000..f892497b252a891974fec5145eb38777920f9fd0 --- /dev/null +++ b/personalised/code/framework/motion_diffusion/losses/losses.py @@ -0,0 +1,62 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..161e39943ae3449b349a53181f9a395443e29ac4 --- /dev/null +++ b/personalised/code/framework/motion_transvae/BasicBlock.py @@ -0,0 +1,109 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..60e56163f767f54ce72718624689602583fde1fa --- /dev/null +++ b/personalised/code/framework/motion_transvae/TransformerVAE.py @@ -0,0 +1,606 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/personalised/code/framework/perfrdiff_rewrite_weight/__init__.py b/personalised/code/framework/perfrdiff_rewrite_weight/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2123b6ca52aa758e224c1db587d25c20c9f8219c --- /dev/null +++ b/personalised/code/framework/perfrdiff_rewrite_weight/__init__.py @@ -0,0 +1,2 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..8c16d6c2c90daa0aa8e3a6ff0ca768a8728fff05 --- /dev/null +++ b/personalised/code/framework/perfrdiff_rewrite_weight/losses.py @@ -0,0 +1,91 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..960da7aed698d14d6a381265c61127e3f2ed1298 --- /dev/null +++ b/personalised/code/framework/perfrdiff_rewrite_weight/modifier/__init__.py @@ -0,0 +1 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..84cae73b2373b0d208bceb1835d94dc3f2ea3ab1 --- /dev/null +++ b/personalised/code/framework/perfrdiff_rewrite_weight/modifier/network.py @@ -0,0 +1,426 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..68460ec32f2d4ef6db56a3cd6f02bdb9ee473031 --- /dev/null +++ b/personalised/code/framework/perfrdiff_rewrite_weight/person_specific/PersonSpecificEncoder.py @@ -0,0 +1,83 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..aa0217d1d639ad48f1b2eeab9ae00120cd71cccf --- /dev/null +++ b/personalised/code/framework/perfrdiff_rewrite_weight/person_specific/__init__.py @@ -0,0 +1 @@ +from .PersonSpecificEncoder import Transformer diff --git a/personalised/code/framework/utils/baseline_reaction_metrics.py b/personalised/code/framework/utils/baseline_reaction_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..5a9edaa39015d0837e0feaa86deb4e8b56aceddf --- /dev/null +++ b/personalised/code/framework/utils/baseline_reaction_metrics.py @@ -0,0 +1,407 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..9ccbd176117c3b7c58fafcfdd544f45db02894d5 --- /dev/null +++ b/personalised/code/framework/utils/compute_metrics.py @@ -0,0 +1,271 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..521f26361a4c73687c90463dfd9962dfdafa88d6 --- /dev/null +++ b/personalised/code/framework/utils/losses.py @@ -0,0 +1,298 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..67a71c656ac2596d81cec4dea16da395ec9ac8e1 --- /dev/null +++ b/personalised/code/framework/utils/losses_causal.py @@ -0,0 +1,52 @@ +""" +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 new file mode 100644 index 0000000000000000000000000000000000000000..4443d42406a9a2c627c34f7299a355966c47be97 --- /dev/null +++ b/personalised/code/framework/utils/util.py @@ -0,0 +1,139 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/personalised/code/launch/blender.py b/personalised/code/launch/blender.py new file mode 100644 index 0000000000000000000000000000000000000000..4ccf32cd715bbd2c8e355f9ccbc39b13f8481872 --- /dev/null +++ b/personalised/code/launch/blender.py @@ -0,0 +1,27 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..617558ed2a849598d4afa20b08442d94fca02794 --- /dev/null +++ b/personalised/code/launch/prepare.py @@ -0,0 +1,75 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..ce82ba7ae21a90fbe5bee1834f692ef1c98eb88d --- /dev/null +++ b/personalised/code/launch/tools.py @@ -0,0 +1,7 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..1e217823bc8f68699ed0b690300c02bf20614300 --- /dev/null +++ b/personalised/code/main.py @@ -0,0 +1,50 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..a5a9e427f680e66be64ab84264b871eb5c44e1f0 --- /dev/null +++ b/personalised/code/requirements.txt @@ -0,0 +1,141 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/personalised/code/trainer/_archive.py b/personalised/code/trainer/_archive.py new file mode 100644 index 0000000000000000000000000000000000000000..94c5b67304d1d9d78d247e0cb8b10dfe249423ee --- /dev/null +++ b/personalised/code/trainer/_archive.py @@ -0,0 +1,255 @@ +# 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 new file mode 100644 index 0000000000000000000000000000000000000000..5063b16293116e530b64f26c04f1328863d41d6e --- /dev/null +++ b/personalised/code/trainer/emotion_autoencoder.py @@ -0,0 +1,338 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..68161d747b3e4a12e172bb207b3b7fdb0902a7ce --- /dev/null +++ b/personalised/code/trainer/g2p_delta.py @@ -0,0 +1,408 @@ +"""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 new file mode 100644 index 0000000000000000000000000000000000000000..bad0d97fea0631f341646f977dfd8919217bebbc --- /dev/null +++ b/personalised/code/trainer/motion_diffusion.py @@ -0,0 +1,1031 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..501fe68455c227059a643ec97121051e6e6a6f8c --- /dev/null +++ b/personalised/code/trainer/motion_diffusion_dp.py @@ -0,0 +1,119 @@ +""" +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 new file mode 100644 index 0000000000000000000000000000000000000000..df9b20fd8ed9d8e4ebcad41726cde555c6bc316a --- /dev/null +++ b/personalised/code/trainer/motion_transvae.py @@ -0,0 +1,788 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..18cfecb35cfc7f30d1d9a2e1ed226d3d99f10713 --- /dev/null +++ b/personalised/code/trainer/perfrdiff_rewrite_weight.py @@ -0,0 +1,1049 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..7c9e6f74e6ce9c7267e004474dbad0dfe91057c7 --- /dev/null +++ b/personalised/code/utils/logging.py @@ -0,0 +1,40 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..81344fc3e6fb63cbf9cdb2fe4330119ed107fce1 --- /dev/null +++ b/personalised/code/utils/render.py @@ -0,0 +1,202 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..cbfb68b92907684e6834e2677cd08c2b297aeac2 --- /dev/null +++ b/personalised/code/utils/runid.py @@ -0,0 +1,18 @@ +""" +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 new file mode 100644 index 0000000000000000000000000000000000000000..ef1b155b477ad768357189419a69102c19f2ff97 --- /dev/null +++ b/personalised/code/utils/util.py @@ -0,0 +1,220 @@ +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 diff --git a/personalised/offline/README.md b/personalised/offline/README.md new file mode 100644 index 0000000000000000000000000000000000000000..28217c0d49f017681e2e767d0676d3907ffa95ab --- /dev/null +++ b/personalised/offline/README.md @@ -0,0 +1,96 @@ +# CaReDiff Personalised Models, Offline Track (REACT 2026) + +Three personalised models for the offline MAFRG track. Each model is the same +frozen generic offline backbone plus a Personalised Residual Adapter (PRA) +trained under a different listener condition. The backbone weights are shared +by all three and are identical to the generic offline submission. + +## Layout + +``` +offline/ + backbone/ frozen generic backbone (shared by all conditions) + CausalTransformerDenoiser/checkpoint_120.pth + DiffusionPriorNetwork/checkpoint_120.pth + EEGPredictionHead/checkpoint_120.pth + adapters/ + personality/ModifierNetwork/checkpoint_best.pth + lhfb/ModifierNetwork/checkpoint_best.pth + both/ModifierNetwork/checkpoint_best.pth +``` + +The adapter file also contains the fine-tuned EEG head, which overwrites the +backbone EEG head at load time. + +## Checksums (SHA-256) + +| File | SHA-256 | +|---|---| +| backbone/CausalTransformerDenoiser/checkpoint_120.pth | 68faca9700415c949eecbe7bd3e381877a76b5e1b24bdab9c30e6fd5b628faa2 | +| backbone/DiffusionPriorNetwork/checkpoint_120.pth | d1b66e87f51afd9bb93bdcef1b9e350e6366aa8f995920e400d7e7dd4e299357 | +| backbone/EEGPredictionHead/checkpoint_120.pth | 750c49999a180cda330b88d771f99d1dca0fd94a810470ea77a45561cfd58780 | +| adapters/personality/ModifierNetwork/checkpoint_best.pth | 8e0a501237c9b80b8c9e9524bd089fa5ca54ad747bdf9ed65dd97b8d883bf928 | +| adapters/lhfb/ModifierNetwork/checkpoint_best.pth | 0ddfde5284c580c3cc461006b2b7cd4df73d2715a8700d3838d8d5e5db8eb7f4 | +| adapters/both/ModifierNetwork/checkpoint_best.pth | 73434669c633bc6384acc9845e62c0c4302c9322be27f04e30005e55dda3ab92 | + +## Conditions + +| Folder | Listener condition | Config value | +|---|---|---| +| adapters/personality | Big-Five personality (5-d) | `personality_only` | +| adapters/lhfb | Listener historical facial behaviour (3DMM) | `3dmm_only` | +| adapters/both | Both, gated fusion | `3dmm_personality` | + +Training: AdamW, learning rate 2e-4, weight decay 1e-4, gradient clipping 1.0, +30 epochs, batch size 32, seed 1234, counterfactual listener-swap loss +(weight 0.5, margin 0.05). The backbone stays frozen throughout. + +## Test performance (MARS test set, official evaluation code, num_gts=10) + +| Condition | FRCorr | FRDist | FRDiv | FRVar | FRRea | FRSyn | +|---|---|---|---|---|---|---| +| personality | 0.7786 | 173.63 | 0.1221 | 0.0782 | 50.91 | 48.37 | +| lhfb | 0.7824 | 173.11 | 0.1200 | 0.0766 | 51.23 | 48.26 | +| both | 0.7822 | 171.41 | 0.1187 | 0.0761 | 50.82 | 48.28 | + +FRRea is the FID between rendered generated frames and ground-truth frames +(56,100 frames per side, frame stride 30). + +## How to run + +The shared source code is in `../code/`. Example for the personality +condition (set `PKG` to the absolute path of the `personalised` folder): + +```bash +cd code +python main.py --config-name g2p_delta stage=test task=offline \ + data_dir= run_id=eval_offline_personality \ + trainer.batch_size=4 num_gts=10 \ + trainer.generic.eval_condition_mode=matched \ + trainer.generic.eval_eeg=false \ + trainer.main_model.args.personal_condition_mode=personality_only \ + resume_id=personality \ + trainer.ckpt_dir=$PKG/offline/adapters \ + trainer.pretrained.diffusion_decoder=$PKG/offline/backbone/CausalTransformerDenoiser/checkpoint_120.pth \ + trainer.pretrained.diffusion_prior=$PKG/offline/backbone/DiffusionPriorNetwork/checkpoint_120.pth \ + trainer.pretrained.eeg_head_checkpoint=$PKG/offline/backbone/EEGPredictionHead/checkpoint_120.pth +``` + +The adapter is loaded from `//ModifierNetwork/`, +which maps directly onto the `adapters/` layout above. For the other two +conditions, change `personal_condition_mode` and `resume_id` (`lhfb` or +`both`) according to the table. The loader verifies that the checkpoint was +trained with the configured condition mode and stops with an error on a +mismatch. + +## Notes + +- Large assets shared with the official baseline are not duplicated here. + The post-processor EmotionVAE checkpoint (517 MB) is required for + evaluation and must be placed at + `code/pretrained_models/post_processor/checkpoint.pth`. The PIRender + renderer (234 MB) is needed only for FRRea rendering. Take both from the + official baseline_react2026 repository. +- Python dependencies: `code/requirements.txt`. +- The MARS dataset is not included and must be obtained through the + challenge organisers. diff --git a/personalised/offline/adapters/both/ModifierNetwork/checkpoint_best.pth b/personalised/offline/adapters/both/ModifierNetwork/checkpoint_best.pth new file mode 100644 index 0000000000000000000000000000000000000000..972b8fbe76f240b68f4945978495b83babf1ac6f --- /dev/null +++ b/personalised/offline/adapters/both/ModifierNetwork/checkpoint_best.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:73434669c633bc6384acc9845e62c0c4302c9322be27f04e30005e55dda3ab92 +size 12226442 diff --git a/personalised/offline/adapters/lhfb/ModifierNetwork/checkpoint_best.pth b/personalised/offline/adapters/lhfb/ModifierNetwork/checkpoint_best.pth new file mode 100644 index 0000000000000000000000000000000000000000..97b50b3ec9540e73898fa86e30818e03e4903251 --- /dev/null +++ b/personalised/offline/adapters/lhfb/ModifierNetwork/checkpoint_best.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0ddfde5284c580c3cc461006b2b7cd4df73d2715a8700d3838d8d5e5db8eb7f4 +size 9420474 diff --git a/personalised/offline/adapters/personality/ModifierNetwork/checkpoint_best.pth b/personalised/offline/adapters/personality/ModifierNetwork/checkpoint_best.pth new file mode 100644 index 0000000000000000000000000000000000000000..51bf872eee3a49bb65a9770ab649a72eeb04d26b --- /dev/null +++ b/personalised/offline/adapters/personality/ModifierNetwork/checkpoint_best.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8e0a501237c9b80b8c9e9524bd089fa5ca54ad747bdf9ed65dd97b8d883bf928 +size 6804442 diff --git a/personalised/offline/backbone/CausalTransformerDenoiser/checkpoint_120.pth b/personalised/offline/backbone/CausalTransformerDenoiser/checkpoint_120.pth new file mode 100644 index 0000000000000000000000000000000000000000..db073eaeb269b93d82c35a46fffb6bb6f8a42630 --- /dev/null +++ b/personalised/offline/backbone/CausalTransformerDenoiser/checkpoint_120.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:68faca9700415c949eecbe7bd3e381877a76b5e1b24bdab9c30e6fd5b628faa2 +size 561047170 diff --git a/personalised/offline/backbone/DiffusionPriorNetwork/checkpoint_120.pth b/personalised/offline/backbone/DiffusionPriorNetwork/checkpoint_120.pth new file mode 100644 index 0000000000000000000000000000000000000000..8f1ad5eb4ce0e6c41b4dd739dc796eb68a04cc76 --- /dev/null +++ b/personalised/offline/backbone/DiffusionPriorNetwork/checkpoint_120.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d1b66e87f51afd9bb93bdcef1b9e350e6366aa8f995920e400d7e7dd4e299357 +size 481319450 diff --git a/personalised/offline/backbone/EEGPredictionHead/checkpoint_120.pth b/personalised/offline/backbone/EEGPredictionHead/checkpoint_120.pth new file mode 100644 index 0000000000000000000000000000000000000000..3535b3e8ad5d04d5fccf900d8cfe47d02456cea4 --- /dev/null +++ b/personalised/offline/backbone/EEGPredictionHead/checkpoint_120.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:750c49999a180cda330b88d771f99d1dca0fd94a810470ea77a45561cfd58780 +size 418238964 diff --git a/personalised/online/README.md b/personalised/online/README.md new file mode 100644 index 0000000000000000000000000000000000000000..e5bb8dfd1f348c34742c968e89d775827dc6fa1a --- /dev/null +++ b/personalised/online/README.md @@ -0,0 +1,106 @@ +# CaReDiff Personalised Models, Online Track (REACT 2026) + +Three personalised models for the online MAFRG track. Each model is the same +frozen generic online backbone plus a Personalised Residual Adapter (PRA) +trained under a different listener condition. Each online adapter was +warm-started from its offline counterpart of the same condition and adapted +with scheduled sampling (probability ramping to 0.5 over 25 epochs). The +backbone weights are shared by all three and are identical to the generic +online submission. + +## Layout + +``` +online/ + backbone/ frozen generic backbone (shared by all conditions) + CausalTransformerDenoiser/checkpoint_120.pth + DiffusionPriorNetwork/checkpoint_120.pth + EEGPredictionHead/checkpoint_120.pth + adapters/ + personality/ModifierNetwork/checkpoint_best.pth + lhfb/ModifierNetwork/checkpoint_best.pth + both/ModifierNetwork/checkpoint_best.pth +``` + +The adapter file also contains the fine-tuned EEG head, which overwrites the +backbone EEG head at load time. + +## Checksums (SHA-256) + +| File | SHA-256 | +|---|---| +| backbone/CausalTransformerDenoiser/checkpoint_120.pth | f4fc53506fc94a65e86b52bfe1491669a73ca3429ad0b1ab51c62488854242f0 | +| backbone/DiffusionPriorNetwork/checkpoint_120.pth | 8b717d619cd37fc793f80f37a4af607bda5e9709c83b82f17916b2467a4380a6 | +| backbone/EEGPredictionHead/checkpoint_120.pth | 60c7a7ae4e6a233fdb59c0ee1e099daf1158931d876a5f46386a781fa2a52a52 | +| adapters/personality/ModifierNetwork/checkpoint_best.pth | fad7691aea8c1895a11fcc7d40873d83ea44e0692e6dfb75f736d89cdd9e62d0 | +| adapters/lhfb/ModifierNetwork/checkpoint_best.pth | c5d8c774774b8a1994be3149c3bb475384326e24a9af64edbf15cda0292da868 | +| adapters/both/ModifierNetwork/checkpoint_best.pth | 77a744f486c46484dc7a357c484bd5c3582345ea5f6947aea1aa625ba00e660a | + +## Conditions + +| Folder | Listener condition | Config value | +|---|---|---| +| adapters/personality | Big-Five personality (5-d) | `personality_only` | +| adapters/lhfb | Listener historical facial behaviour (3DMM) | `3dmm_only` | +| adapters/both | Both, gated fusion | `3dmm_personality` | + +Training: AdamW, learning rate 2e-4, weight decay 1e-4, gradient clipping 1.0, +30 epochs, batch size 32, seed 1234, counterfactual listener-swap loss +(weight 0.5, margin 0.05). The backbone stays frozen throughout. Generation +is autoregressive over 30-frame windows; each of the 10 predictions +conditions on its own previous window. + +## Test performance (MARS test set, official evaluation code, num_gts=10) + +| Condition | FRCorr | FRDist | FRDiv | FRVar | FRRea | FRSyn | +|---|---|---|---|---|---|---| +| personality | 0.6485 | 185.17 | 0.1521 | 0.0831 | 50.58 | 47.92 | +| lhfb | 0.6481 | 191.88 | 0.1521 | 0.0828 | 50.89 | 47.92 | +| both | 0.6355 | 181.11 | 0.1451 | 0.0790 | 52.09 | 48.16 | + +FRRea is the FID between rendered generated frames and ground-truth frames +(56,100 frames per side, frame stride 30). + +## How to run + +The shared source code is in `../code/`. Example for the personality +condition (set `PKG` to the absolute path of the `personalised` folder): + +```bash +cd code +python main.py --config-name g2p_delta_online stage=test task=online \ + data_dir= run_id=eval_online_personality \ + trainer.batch_size=4 num_gts=10 \ + trainer.generic.eval_condition_mode=matched \ + trainer.generic.eval_eeg=false \ + trainer.main_model.args.personal_condition_mode=personality_only \ + trainer.model.diff_model.diffusion_decoder.args.past_l_emotion_drop_prob=0.2 \ + resume_id=personality \ + trainer.ckpt_dir=$PKG/online/adapters \ + trainer.pretrained.diffusion_decoder=$PKG/online/backbone/CausalTransformerDenoiser/checkpoint_120.pth \ + trainer.pretrained.diffusion_prior=$PKG/online/backbone/DiffusionPriorNetwork/checkpoint_120.pth \ + trainer.pretrained.eeg_head_checkpoint=$PKG/online/backbone/EEGPredictionHead/checkpoint_120.pth +``` + +The adapter is loaded from `//ModifierNetwork/`, +which maps directly onto the `adapters/` layout above. For the other two +conditions, change `personal_condition_mode` and `resume_id` (`lhfb` or +`both`) according to the table. The loader verifies that the checkpoint was +trained with the configured condition mode and stops with an error on a +mismatch. + +Keep `past_l_emotion_drop_prob=0.2`. This flag enables the past-listener +conditioning pathway; the config default disables it and does not reproduce +the reported numbers. + +## Notes + +- Large assets shared with the official baseline are not duplicated here. + The post-processor EmotionVAE checkpoint (517 MB) is required for + evaluation and must be placed at + `code/pretrained_models/post_processor/checkpoint.pth`. The PIRender + renderer (234 MB) is needed only for FRRea rendering. Take both from the + official baseline_react2026 repository. +- Python dependencies: `code/requirements.txt`. +- The MARS dataset is not included and must be obtained through the + challenge organisers. diff --git a/personalised/online/adapters/both/ModifierNetwork/checkpoint_best.pth b/personalised/online/adapters/both/ModifierNetwork/checkpoint_best.pth new file mode 100644 index 0000000000000000000000000000000000000000..f4803711303cf6e05f14d656e6787ea02d819ca0 --- /dev/null +++ b/personalised/online/adapters/both/ModifierNetwork/checkpoint_best.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:77a744f486c46484dc7a357c484bd5c3582345ea5f6947aea1aa625ba00e660a +size 12226442 diff --git a/personalised/online/adapters/lhfb/ModifierNetwork/checkpoint_best.pth b/personalised/online/adapters/lhfb/ModifierNetwork/checkpoint_best.pth new file mode 100644 index 0000000000000000000000000000000000000000..9caf7e6d235e4d89ad4f2899fe1c1c1dea73ec87 --- /dev/null +++ b/personalised/online/adapters/lhfb/ModifierNetwork/checkpoint_best.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c5d8c774774b8a1994be3149c3bb475384326e24a9af64edbf15cda0292da868 +size 9420474 diff --git a/personalised/online/adapters/personality/ModifierNetwork/checkpoint_best.pth b/personalised/online/adapters/personality/ModifierNetwork/checkpoint_best.pth new file mode 100644 index 0000000000000000000000000000000000000000..963961f64a10ae6abc658e51ce0f528a05f6aff1 --- /dev/null +++ b/personalised/online/adapters/personality/ModifierNetwork/checkpoint_best.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fad7691aea8c1895a11fcc7d40873d83ea44e0692e6dfb75f736d89cdd9e62d0 +size 6804442 diff --git a/personalised/online/backbone/CausalTransformerDenoiser/checkpoint_120.pth b/personalised/online/backbone/CausalTransformerDenoiser/checkpoint_120.pth new file mode 100644 index 0000000000000000000000000000000000000000..55c7c0b534b1ced68072ae4afc70eec49b3a9f7c --- /dev/null +++ b/personalised/online/backbone/CausalTransformerDenoiser/checkpoint_120.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f4fc53506fc94a65e86b52bfe1491669a73ca3429ad0b1ab51c62488854242f0 +size 561047170 diff --git a/personalised/online/backbone/DiffusionPriorNetwork/checkpoint_120.pth b/personalised/online/backbone/DiffusionPriorNetwork/checkpoint_120.pth new file mode 100644 index 0000000000000000000000000000000000000000..fc3c7ae8445e0c35122793f2bae423d65ad074e8 --- /dev/null +++ b/personalised/online/backbone/DiffusionPriorNetwork/checkpoint_120.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8b717d619cd37fc793f80f37a4af607bda5e9709c83b82f17916b2467a4380a6 +size 481319450 diff --git a/personalised/online/backbone/EEGPredictionHead/checkpoint_120.pth b/personalised/online/backbone/EEGPredictionHead/checkpoint_120.pth new file mode 100644 index 0000000000000000000000000000000000000000..974ce81c6195e186b156b577682d21ea1bc65c70 --- /dev/null +++ b/personalised/online/backbone/EEGPredictionHead/checkpoint_120.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:60c7a7ae4e6a233fdb59c0ee1e099daf1158931d876a5f46386a781fa2a52a52 +size 418238964