Initial release on Huggingface Spaces with Gradio UI

.gitattributes

@@ -33,3 +33,14 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.torchscript filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+./magick filter=lfs diff=lfs merge=lfs -text
+models/* filter=lfs diff=lfs merge=lfs -text
+models/nlf_l_multi.torchscript filter=lfs diff=lfs merge=lfs -text
+models/checkpoint_150.pt filter=lfs diff=lfs merge=lfs -text
+downloads/* filter=lfs diff=lfs merge=lfs -text
+outputs/* filter=lfs diff=lfs merge=lfs -text
+intermediate_results/* filter=lfs diff=lfs merge=lfs -text
+predictions/* filter=lfs diff=lfs merge=lfs -text
+predictions/joints3d.npy filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,213 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual Environment
+venv/
+ENV/
+env/
+.env
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Project specific
+*.pth
+*.ckpt
+# *.pt
+*.bin
+*.npy
+*.npz
+*.mp4
+*.avi
+*.mov
+*.jpg
+*.jpeg
+*.png
+
+# Logs
+*.log
+logs/
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+# downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# PyPI configuration file
+.pypirc
+9622_GRAB/
+magick
+outputs/*_obj
+outputs/
+intermediate_results/

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

.vscode/settings.json

@@ -0,0 +1,6 @@
+{
+    "python.analysis.extraPaths": [
+        "./SkeletonDiffusion/src",
+        "./src_joints2smpl_demo/convert_"
+    ]
+}

README.md

@@ -4,10 +4,12 @@ emoji: 💻
 colorFrom: purple
 colorTo: green
 sdk: gradio
-sdk_version: 5.12.0
+sdk_version: 5.24.0
 app_file: app.py
 pinned: false
 license: mit
+run: |
+  bash setup.sh
+  python app.py
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

SkeletonDiffusion/configs/config_eval/config.yaml

@@ -0,0 +1,53 @@
+hydra:
+  output_subdir: null  
+  run:  
+    dir: .  
+  job:
+    chdir: False
+
+dataset_main_path: ./datasets
+dataset_annotation_path: ${dataset_main_path}/annotations&interm
+dataset_precomputed_path: ${dataset_main_path}/processed
+checkpoint_path: ''
+defaults:
+  - _self_
+  - task: hmp
+  - method_specs: skeleton_diffusion 
+  - dataset: amass
+  - override hydra/hydra_logging: disabled  
+  - override hydra/job_logging: disabled  
+
+method_name: ${method_specs.method_name}
+dtype: float32
+if_noisy_obs: False
+noise_level: 0.25
+noise_std: 0.02
+# num_nodes: ${eval:"int(${dataset.num_joints})-int(not ${task.if_consider_hip})"}
+
+stats_mode: deterministic #probabilistic, deterministic
+batch_size:  512 
+metrics_at_cpu: False
+n_gpu: 1
+num_samples: 50 
+if_measure_time: False
+
+seed: 0
+dataset_split: test
+silent: False
+
+obs_length: ${eval:'int(${task.history_sec} * ${dataset.fps})'}
+pred_length:  ${eval:'int(${task.prediction_horizon_sec} * ${dataset.fps})'}
+
+if_store_output: False
+store_output_path: ${eval:"'models/final_predictions_storage/${task.task_name}/${method_specs.method_name}/${dataset.dataset_name}/' if not ${if_long_term_test} else 'models/final_predictions_storage/${task.task_name}_longterm/${method_specs.method_name}/${dataset.dataset_name}/'"} 
+if_store_gt: False 
+store_gt_path: ${eval:"'models/final_predictions_storage/${task.task_name}/GT/${dataset.dataset_name}/' if not ${if_long_term_test} else 'models/final_predictions_storage/${task.task_name}_longterm/GT/${dataset.dataset_name}/'"} 
+
+if_compute_apde: ${eval:"False if ${eval:"'${dataset.dataset_name}' in ['freeman', '3dpw', 'nymeria']"} else True"}
+if_long_term_test: False
+long_term_factor: 2.5
+
+if_compute_fid: False
+if_compute_cmd: False
+
+

SkeletonDiffusion/configs/config_eval/config_inferencetime.yaml

@@ -0,0 +1,43 @@
+hydra:
+  output_subdir: null  
+  run:  
+    dir: .  
+  job:
+    chdir: False
+
+dataset_main_path: ./datasets
+dataset_annotation_path: ${dataset_main_path}/annotations&interm
+dataset_precomputed_path: ${dataset_main_path}/processed
+checkpoint_path: ''
+defaults:
+  - _self_
+  - task: hmp
+  - method_specs: skeldiff 
+  - dataset: amass
+  - override hydra/hydra_logging: disabled  
+  - override hydra/job_logging: disabled  
+
+method_name: ${method_specs.method_name}
+dtype: float32
+if_noisy_obs: False
+noise_level: 0.25
+noise_std: 0.02
+
+mode: stats # 'vis: visualize results\ngen: generate and store all visualizations for a single batch\nstats: launch numeric evaluation')
+stats_mode: deterministic #
+if_measure_time: True
+batch_size:  1
+metrics_at_cpu: False
+n_gpu: 1
+num_samples: 50 
+
+seed: 0
+dataset_split: test
+silent: False
+
+obs_length: ${eval:'int(${task.history_sec} * ${dataset.fps})'}
+pred_length:  ${eval:'int(${task.prediction_horizon_sec} * ${dataset.fps})'}
+
+if_long_term_test: False
+
+

SkeletonDiffusion/configs/config_eval/dataset/3dpw.yaml

@@ -0,0 +1,35 @@
+num_joints: 22 #including the hip root joint
+fps: 60
+
+multimodal_threshold: 0.4
+dataset_type: D3PWZeroShotDataset
+dataset_name: 3dpw
+precomputed_folder: "${dataset_precomputed_path}/3DPW/${task.task_name}/"
+annotations_folder: "${dataset_annotation_path}/3DPW/${task.task_name}/"
+dtype: float32
+
+data_loader_train_eval:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    da_mirroring: 0.
+    da_rotations: 0.
+    drop_last: False
+    if_load_mmgt: False
+
+data_loader_valid:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    segments_path: "${dataset.annotations_folder}/segments_valid.csv"
+    actions: "all"
+    drop_last: False
+    if_load_mmgt: False
+
+
+data_loader_test:
+    shuffle: False
+    segments_path: "${dataset.annotations_folder}/segments_test_zero_shot.csv"
+    actions: "all"
+    drop_last: False
+    if_load_mmgt: ${eval:'True if "probabilistic" in "${stats_mode}" else False'}

SkeletonDiffusion/configs/config_eval/dataset/amass-mano.yaml

@@ -0,0 +1,76 @@
+num_joints: 52 #included hip even if if_consider_hip=False
+fps: 60
+
+multimodal_threshold: 0.4
+dataset_type: AMASSDataset
+dataset_name: amass-mano
+precomputed_folder: "${dataset_precomputed_path}/AMASS-MANO/${task.task_name}/"
+annotations_folder: "${dataset_annotation_path}/AMASS-MANO/${task.task_name}/"
+dtype: float32
+
+# Accordingly, the training set
+# contains the ACCAD, BMLhandball, BMLmovi, BMLrub,
+# CMU, EKUT, EyesJapanDataset, KIT, PosePrior, TCD-
+# Hands, and TotalCapture datasets, and the validation set
+# contains the HumanEva, HDM05, SFU, and MoSh datasets.
+# The remaining datasets are all part of the test set: DFaust,
+# DanceDB, GRAB, HUMAN4D, SOMA, SSM, and Transi-
+# tions.
+
+
+data_loader_train:
+    stride: 60
+    augmentation: 30
+    shuffle: True
+    datasets: ['ACCAD', "BMLhandball", "BMLmovi", "BMLrub", 'EKUT', 'CMU', 'EyesJapanDataset', 'KIT', "PosePrior", 'TCDHands', 'TotalCapture'] # from paper
+               
+                
+                
+                # "EyesJapanDataset",
+                
+                # ,
+                
+                # "HDM05",
+                # "MoSh"
+    da_mirroring: 0.5
+    da_rotations: 1.0
+    drop_last: True
+    if_load_mmgt: False
+
+
+data_loader_train_eval:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    datasets: ['ACCAD', "BMLhandball", "BMLmovi", "BMLrub", 'EKUT', 'CMU', 'EyesJapanDataset', 'KIT', "PosePrior", 'TCDHands', 'TotalCapture'] # from paper
+    # datasets: ['ACCAD', "BMLhandball", "BMLmovi",  'CMU', 'KIT', 'TotalCapture'] # decrease evaluatio time
+    da_mirroring: 0.
+    da_rotations: 0.
+    drop_last: False
+    if_load_mmgt: False
+
+data_loader_valid:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    # segments_path: "./dataset_annotation_path/FreeMan/${task.task_name}/segments_valid.csv"
+    datasets: ['HumanEva', 'HDM05', 'SFU', 'MoSh'] # from paper
+    file_idces: "all"
+    drop_last: False
+    if_load_mmgt: False
+
+
+data_loader_test:
+    shuffle: False
+    segments_path: "${dataset.annotations_folder}/segments_test.csv"
+    # datasets:  ['Transitions_mocap', 'SSM_synced'],     #DFaust, DanceDB, GRAB, HUMAN4D, SOMA, SSM, and Transitions.
+    datasets:  
+        - Transitions 
+        - SSM 
+        - DFaust 
+        - DanceDB 
+        - GRAB 
+        - HUMAN4D 
+        - SOMA 
+    drop_last: False
+    if_load_mmgt: False

SkeletonDiffusion/configs/config_eval/dataset/amass.yaml

@@ -0,0 +1,52 @@
+num_joints: 22 #including the hip root joint
+fps: 60
+
+multimodal_threshold: 0.4
+dataset_type: AMASSDataset
+dataset_name: "amass"
+precomputed_folder: "${dataset_precomputed_path}/AMASS/${task.task_name}/"
+annotations_folder: "${dataset_annotation_path}/AMASS/${task.task_name}/"
+dtype: float32
+
+# Accordingly, the training set
+# contains the ACCAD, BMLhandball, BMLmovi, BMLrub,
+# CMU, EKUT, EyesJapanDataset, KIT, PosePrior, TCD-
+# Hands, and TotalCapture datasets, and the validation set
+# contains the HumanEva, HDM05, SFU, and MoSh datasets.
+# The remaining datasets are all part of the test set: DFaust,
+# DanceDB, GRAB, HUMAN4D, SOMA, SSM, and Transi-
+# tions.
+
+data_loader_train_eval:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    datasets: ['ACCAD', "BMLhandball", "BMLmovi", "BMLrub", 'EKUT', 'CMU', 'EyesJapanDataset', 'KIT', "PosePrior", 'TCDHands', 'TotalCapture']
+    da_mirroring: 0.
+    da_rotations: 0.
+    drop_last: False
+    if_load_mmgt: False
+
+data_loader_valid:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    datasets: ['HumanEva', 'HDM05',  'SFU',  MoSh']
+    file_idces: "all"
+    drop_last: False
+    if_load_mmgt: False
+
+
+data_loader_test:
+    shuffle: False
+    segments_path: ${eval:"'${dataset.annotations_folder}/segments_test.csv' if not ${if_long_term_test}  else '${dataset.annotations_folder}/segments_5s_test_long_term_pred.csv'"}
+    datasets:  
+        - Transitions 
+        - SSM 
+        - DFaust
+        - DanceDB 
+        - GRAB 
+        - HUMAN4D 
+        - SOMA 
+    drop_last: False
+    if_load_mmgt: ${eval:'True if "probabilistic" in "${stats_mode}" else False'}

SkeletonDiffusion/configs/config_eval/dataset/freeman.yaml

@@ -0,0 +1,23 @@
+num_joints: 18 #including the hip root joint
+fps: 30
+dataset_type: FreeManDataset
+dataset_name: "freeman"
+precomputed_folder: "${dataset_precomputed_path}/FreeMan/${task.task_name}/"
+annotations_folder: "${dataset_annotation_path}/FreeMan/${task.task_name}/"
+dtype: float32
+multimodal_threshold: 0.5
+
+data_loader_valid:
+    shuffle: False
+    segments_path: "${dataset.annotations_folder}/segments_valid.csv"
+    actions: "all"
+    drop_last: False
+    if_load_mmgt: ${eval:'True if "probabilistic in str(${stats_mode})" else False'}
+
+
+data_loader_test:
+    shuffle: False
+    segments_path: "${dataset.annotations_folder}/segments_test.csv"
+    actions: "all"
+    drop_last: False
+    if_load_mmgt: ${eval:'True if "probabilistic" in "${stats_mode}" else False'}

SkeletonDiffusion/configs/config_eval/dataset/h36m.yaml

@@ -0,0 +1,26 @@
+num_joints: 17 #including the hip root joint
+fps: 50
+dataset_type: H36MDataset
+dataset_name: "h36m"
+precomputed_folder: "${dataset_precomputed_path}/Human36M/${task.task_name}/"
+annotations_folder: "${dataset_annotation_path}/Human36M/${task.task_name}"
+dtype: float32
+multimodal_threshold: 0.5
+
+data_loader_valid:
+    augmentation: 0
+    shuffle: False
+    subjects: ["S8"]
+    segments_path: "${dataset.annotations_folder}/segments_valid.csv"
+    actions: "all"
+    drop_last: False
+    if_load_mmgt:  ${eval:'True if "probabilistic in str(${stats_mode})" else False'}
+
+data_loader_test:
+    shuffle: False
+    augmentation: 0
+    segments_path: "${dataset.annotations_folder}/segments_test.csv"
+    subjects: ["S9", "S11"]
+    actions: "all"
+    drop_last: False
+    if_load_mmgt: ${eval:'True if "probabilistic" in "${stats_mode}" else False'}

SkeletonDiffusion/configs/config_eval/method_specs/skeleton_diffusion.yaml

@@ -0,0 +1 @@
+method_name: SkeletonDiffusion

SkeletonDiffusion/configs/config_eval/method_specs/zerovelocity_alg_baseline.yaml

@@ -0,0 +1,3 @@
+motion_repr_type: "SkeletonCenterPose"
+method_name: ZeroVelocityBaseline
+baseline_out_path: ./models/output/baselines

SkeletonDiffusion/configs/config_eval/task/hmp.yaml

@@ -0,0 +1,4 @@
+history_sec: 0.5 #${eval:'float(1) if ${task}=="motpred" else float(0.5)'}
+prediction_horizon_sec: 2 # ${eval:"float(4) if ${task}=='motpred' else float(2)"}
+task_name: "hmp"
+if_consider_hip: False

SkeletonDiffusion/configs/config_train/config_autoencoder.yaml

@@ -0,0 +1,27 @@
+defaults:
+  - _self_
+  - task: hmp
+  - dataset: h36m
+  - model: autoencoder
+  - override hydra/job_logging: disabled    
+  # - override hydra/hydra_logging: disabled  
+
+dataset_main_path: ./datasets
+dataset_annotation_path: ${dataset_main_path}/annotations&interm
+dataset_precomputed_path: ${dataset_main_path}/processed
+if_resume_training: false
+debug: false
+device: cuda
+load: false
+load_path: ''
+output_log_path: ../../my_exps/output/${task.task_name}/${dataset.dataset_name}/autoencoder/${now:%B%d_%H-%M-%S}_ID${slurm_id}_${info}
+slurm_id: 0
+slurm_first_run: None
+info: ''
+hydra:
+  run:
+    dir: ${output_log_path}
+  job:
+    chdir: False
+
+

SkeletonDiffusion/configs/config_train/dataset/amass.yaml

@@ -0,0 +1,48 @@
+num_joints: 22 #including the hip root joint
+fps: 60
+
+multimodal_threshold: 0.4
+dataset_type: AMASSDataset
+dataset_name: amass
+precomputed_folder: "${dataset_precomputed_path}/AMASS/${task.task_name}/"
+annotations_folder: "${dataset_annotation_path}/AMASS/${task.task_name}/"
+dtype: float32
+
+# Accordingly, the training set
+# contains the ACCAD, BMLhandball, BMLmovi, BMLrub,
+# CMU, EKUT, EyesJapanDataset, KIT, PosePrior, TCD-
+# Hands, and TotalCapture datasets, and the validation set
+# contains the HumanEva, HDM05, SFU, and MoSh datasets.
+# The remaining datasets are all part of the test set: DFaust,
+# DanceDB, GRAB, HUMAN4D, SOMA, SSM, and Transi-
+# tions.
+
+data_loader_train:
+    stride: 60
+    augmentation: 30
+    shuffle: True
+    datasets: ['ACCAD', "BMLhandball", "BMLmovi", "BMLrub", 'EKUT', 'CMU', 'EyesJapanDataset', 'KIT', "PosePrior", 'TCDHands', 'TotalCapture']
+    da_mirroring: 0.5
+    da_rotations: 1.0
+    drop_last: True
+    if_load_mmgt: False
+
+
+data_loader_train_eval:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    datasets: ['ACCAD', "BMLhandball", "BMLmovi", "BMLrub", 'EKUT', 'CMU', 'EyesJapanDataset', 'KIT', "PosePrior", 'TCDHands', 'TotalCapture']
+    da_mirroring: 0.
+    da_rotations: 0.
+    drop_last: False
+    if_load_mmgt: False
+
+data_loader_valid:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    datasets: ['HumanEva', 'HDM05', 'SFU', 'MoSh']
+    file_idces: "all"
+    drop_last: False
+    if_load_mmgt: False

SkeletonDiffusion/configs/config_train/dataset/freeman.yaml

@@ -0,0 +1,38 @@
+num_joints: 18 #including the hip root joint
+fps: 30
+
+multimodal_threshold: 0.5
+dataset_type: FreeManDataset
+dataset_name: freeman
+precomputed_folder: "${dataset_precomputed_path}/FreeMan/${task.task_name}/"
+annotations_folder: "${dataset_annotation_path}/FreeMan/${task.task_name}/"
+dtype: float32
+
+data_loader_train:
+    stride: 10
+    augmentation: 5
+    shuffle: True
+    actions: "all"
+    da_mirroring: 0.5
+    da_rotations: 1.0
+    drop_last: True
+    if_load_mmgt: False
+
+data_loader_train_eval:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    actions: "all"
+    da_mirroring: 0.
+    da_rotations: 0.
+    drop_last: False
+    if_load_mmgt: False
+
+data_loader_valid:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    segments_path: "${dataset.annotations_folder}/segments_valid.csv"
+    actions: "all"
+    drop_last: False
+    if_load_mmgt: False

SkeletonDiffusion/configs/config_train/dataset/h36m.yaml

@@ -0,0 +1,40 @@
+num_joints: 17  #including the hip root joint
+fps: 50
+
+multimodal_threshold: 0.5
+dataset_type: H36MDataset
+dataset_name: h36m
+precomputed_folder: "${dataset_precomputed_path}/Human36M/${task.task_name}/"
+annotations_folder: "${dataset_annotation_path}/Human36M/${task.task_name}"
+dtype: float32
+data_loader_train:
+    stride: 10
+    augmentation: 5
+    shuffle: True
+    subjects: ["S1", "S5", "S6", "S7", "S8"] # training on the validation split as well as in BeLFusion, CoMusion
+    actions: "all"
+    da_mirroring: 0.5
+    da_rotations: 1.0
+    drop_last: True
+    if_load_mmgt: False
+
+data_loader_train_eval:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    subjects: ["S1", "S5", "S6", "S7", "S8"]
+    actions: "all"
+    da_mirroring: 0.
+    da_rotations: 0.
+    drop_last: False
+    if_load_mmgt: False
+
+data_loader_valid:
+    stride: 30
+    augmentation: 0
+    shuffle: False
+    subjects: ["S8"]
+    segments_path: "${dataset.annotations_folder}/segments_valid.csv"
+    actions: "all"
+    drop_last: False
+    if_load_mmgt: False

SkeletonDiffusion/configs/config_train/model/autoencoder.yaml

@@ -0,0 +1,57 @@
+batch_size: 64
+batch_size_eval: 512
+eval_frequency: 5
+num_epochs: 200
+num_iteration_eval: 10
+num_workers: 4
+seed: 52345
+
+use_lr_scheduler: True
+lr_scheduler_kwargs: 
+  lr_scheduler_type: ExponentialLRSchedulerWarmup  
+  warmup_duration: 10
+  update_every: 1
+  min_lr: 1.e-4
+  gamma_decay: 0.98
+
+
+loss_pose_type: l1
+
+
+lr: 0.5e-2
+
+latent_size: 96
+output_size: 3 #128
+
+
+z_activation: tanh
+
+num_iter_perepoch: ${eval:"int(485) if ${eval:"'${dataset.dataset_name}' == 'h36m'"} else 580"}
+
+obs_length: ${eval:'int(${task.history_sec} * ${dataset.fps})'}
+prediction_horizon_train: ${model.prediction_horizon}
+prediction_horizon_eval: ${model.prediction_horizon}
+prediction_horizon:  ${eval:'int(${task.prediction_horizon_sec} * ${dataset.fps})'}
+pred_length: ${model.prediction_horizon_eval}
+
+
+
+autoenc_arch:
+  enc_num_layers: 1
+  encoder_hidden_size: 96
+  decoder_hidden_size: 96
+  arch: AutoEncoder 
+  recurrent_arch_enc: StaticGraphGRU
+  recurrent_arch_decoder: StaticGraphGRU
+
+
+
+prediction_horizon_train_min: 10 
+prediction_horizon_train_min_from_epoch: 200
+curriculum_it: 10
+random_prediction_horizon: True
+
+
+
+
+  

SkeletonDiffusion/configs/config_train/task/hmp.yaml

@@ -0,0 +1,11 @@
+if_consider_hip: False
+
+history_sec: 0.5
+prediction_horizon_sec: 2
+
+
+# Joint representation & Skeleton specs
+motion_repr_type: "SkeletonRescalePose"
+pose_box_size: 1.5 # in meters
+seq_centering: 0
+task_name: hmp

SkeletonDiffusion/configs/config_train_diffusion/config_diffusion.yaml

@@ -0,0 +1,25 @@
+if_resume_training: false
+debug: false
+device: cuda
+load: false
+load_path: ''
+dataset_main_path: ./datasets
+dataset_annotation_path: ${dataset_main_path}/annotations&interm
+dataset_precomputed_path: ${dataset_main_path}/processed
+_load_saved_aoutoenc: hmp-h36m
+
+output_log_path: ../../my_exps/output/${eval:"'${_load_saved_aoutoenc}'.split('-', 1)[0]"}/${eval:"'${_load_saved_aoutoenc}'.split('-', 1)[1]"}/diffusion/${now:%B%d_%H-%M-%S}_ID${slurm_id}_${info}
+slurm_id: 0
+slurm_first_run: None
+info: ''
+hydra:
+  run:
+    dir: ${output_log_path}
+  job:
+    chdir: False
+
+defaults:
+  - _self_
+  - model: skeleton_diffusion
+  - cov_matrix: adjacency
+  - override hydra/job_logging: disabled  

SkeletonDiffusion/configs/config_train_diffusion/cov_matrix/adjacency.yaml

@@ -0,0 +1 @@
+covariance_matrix_type: adjacency

SkeletonDiffusion/configs/config_train_diffusion/cov_matrix/reachability.yaml

@@ -0,0 +1,3 @@
+covariance_matrix_type: reachability
+reachability_matrix_degree_factor: 0.5
+reachability_matrix_stop_at: hips # or None 'hips' or null

SkeletonDiffusion/configs/config_train_diffusion/model/isotropic_diffusion.yaml

@@ -0,0 +1,57 @@
+pretrained_GM_folder: ${eval:"'./models/final_checkpoints/H36M/hmp/autoencoder/January19_19-24-04_ID1137310' if ${eval:"'${_load_saved_aoutoenc}'.split('-')[1] == 'h36m'"}   else './models/final_checkpoints/AMASS/hmp/autoencoder/May11_10-35-09_ID1185354'"} 
+pretrained_autoencoder_path: '${model.pretrained_GM_folder}/checkpoints/checkpoint_final.pt' 
+
+# These options have to still be checked (ablation)
+lr: 1.e-3 
+diffusion_objective: pred_x0
+weight_decay: 0.
+
+
+if_use_ema: True
+step_start_ema: 100 #100 is default
+ema_power: ${eval:'2/3'} #${eval:'3/4'} #
+ema_update_every: 10
+ema_min_value: 0.0
+use_lr_scheduler: True
+lr_scheduler_kwargs: 
+  lr_scheduler_type: ExponentialLRSchedulerWarmup # or SchedulerReduceLROnPlateau
+  warmup_duration: 200
+  update_every: 10
+  min_lr: 1.e-4
+  gamma_decay: 0.98
+
+# THese option are already ablated
+diffusion_conditioning: True
+num_epochs: 600
+num_workers: 4
+batch_size: 64
+batch_size_eval: 256 # or 256 TO CHECK
+eval_frequency: 25 #1000 # in epochs
+train_pick_best_sample_among_k: 50
+similarity_space: latent_space # input_space, latent_space or metric_space
+
+diffusion_activation: identity
+num_prob_samples: 50
+diffusion_timesteps: 10
+
+diffusion_type: IsotropicGaussianDiffusion
+beta_schedule: cosine
+diffusion_loss_type: l1
+num_iter_perepoch: null
+seed: 63485
+
+diffusion_arch: 
+    arch: Denoiser
+    use_attention: True
+    self_condition: False 
+    norm_type: none
+    depth: 1
+    # resnet_block_groups = 8,
+    # learned_variance: False
+    # learned_sinusoidal_cond: False
+    # random_fourier_features: False
+    # learned_sinusoidal_dim: 16
+    # sinusoidal_pos_emb_theta: 10000
+    attn_dim_head: 32
+    attn_heads: 4
+    learn_influence: True

SkeletonDiffusion/configs/config_train_diffusion/model/isotropic_diffusion_in_noniso_class.yaml

@@ -0,0 +1,70 @@
+# pretrained_GM_folder: ./models/final_checkpoints/H36M/hmp/autoencoder/January19_19-24-04_ID1137310
+# _pretrained_GM_checkpoint: checkpoint_final
+# pretrained_autoencoder_path: '${model.pretrained_GM_folder}/checkpoints/${model._pretrained_GM_checkpoint}.pt' 
+
+_pretrained_GM_checkpoint: checkpoint_final
+pretrained_autoencoder_path: '${model.pretrained_GM_folder}/checkpoints/${model._pretrained_GM_checkpoint}.pt' 
+pretrained_GM_folder: ${eval:"'./models/final_checkpoints/H36M/hmp/autoencoder/January19_19-24-04_ID1137310' if ${eval:"'${_load_saved_aoutoenc}'.split('-')[1] == 'h36m'"}   else './models/final_checkpoints/AMASS/hmp/autoencoder/May11_10-35-09_ID1185354'"} 
+
+
+# These options have to still be checked (ablation)
+lr: 1.e-3 
+diffusion_objective: pred_x0
+weight_decay: 0.
+
+
+if_use_ema: True
+step_start_ema: 100 #100 is default
+ema_power: ${eval:'2/3'} #${eval:'3/4'} #
+ema_update_every: 10
+ema_min_value: 0.0
+use_lr_scheduler: True
+lr_scheduler_kwargs: 
+  lr_scheduler_type: ExponentialLRSchedulerWarmup # or SchedulerReduceLROnPlateau
+  warmup_duration: 200
+  update_every: 10
+  min_lr: 1.e-4
+  gamma_decay: 0.98
+
+# THese option are already ablated
+diffusion_conditioning: True
+num_epochs: 800
+num_workers: 4
+batch_size: 64
+batch_size_eval: 256
+eval_frequency: 25 #1000 # in epochs
+train_pick_best_sample_among_k: 50
+similarity_space: latent_space # input_space, latent_space or metric_space
+
+
+diffusion_activation: identity
+num_prob_samples: 50
+diffusion_timesteps: 10
+
+diffusion_type: NonisotropicGaussianDiffusion
+diffusion_loss_type: snr #snr_triangle_inequality,mahalanobis, snr
+loss_reduction_type: l1
+if_run_as_isotropic: True
+if_sigma_n_scale: True
+diffusion_covariance_type: isotropic # anisotropic, isotropic, skeleton-diffusion
+gamma_scheduler: cosine # mono_decrease, cosine
+beta_schedule: cosine
+sigma_n_scale: spectral
+num_iter_perepoch: null
+seed: 63485
+
+diffusion_arch: 
+    arch: Denoiser
+    use_attention: True
+    self_condition: False # True holds better results, but it takes longer to train.
+    norm_type: none
+    depth: 1
+    # resnet_block_groups = 8,
+    # learned_variance: False
+    # learned_sinusoidal_cond: False
+    # random_fourier_features: False
+    # learned_sinusoidal_dim: 16
+    # sinusoidal_pos_emb_theta: 10000
+    attn_dim_head: 32
+    attn_heads: 4
+    learn_influence: True

SkeletonDiffusion/configs/config_train_diffusion/model/skeleton_diffusion.yaml

@@ -0,0 +1,69 @@
+# pretrained_GM_folder: ./models/final_checkpoints/H36M/hmp/autoencoder/January19_19-24-04_ID1137310
+# _pretrained_GM_checkpoint: checkpoint_final
+# pretrained_autoencoder_path: '${model.pretrained_GM_folder}/checkpoints/${model._pretrained_GM_checkpoint}.pt' 
+
+_pretrained_GM_checkpoint: checkpoint_final
+pretrained_autoencoder_path: '${model.pretrained_GM_folder}/checkpoints/${model._pretrained_GM_checkpoint}.pt' 
+pretrained_GM_folder: ${eval:"'./models/final_checkpoints/H36M/hmp/autoencoder/January19_19-24-04_ID1137310' if ${eval:"'${_load_saved_aoutoenc}'.split('-')[1] == 'h36m'"}   else './models/final_checkpoints/AMASS/hmp/autoencoder/May11_10-35-09_ID1185354'"} 
+
+
+# These options have to still be checked (ablation)
+lr: 1.e-3 
+diffusion_objective: pred_x0
+weight_decay: 0.
+
+
+if_use_ema: True
+step_start_ema: 100 #100 is default
+ema_power: ${eval:'2/3'} #${eval:'3/4'} #
+ema_update_every: 10
+ema_min_value: 0.0
+use_lr_scheduler: True
+lr_scheduler_kwargs: 
+  lr_scheduler_type: ExponentialLRSchedulerWarmup # or SchedulerReduceLROnPlateau
+  warmup_duration: 200
+  update_every: 10
+  min_lr: 1.e-4
+  gamma_decay: 0.98
+
+diffusion_conditioning: True
+num_epochs: 800
+num_workers: 4
+batch_size: 64
+batch_size_eval: 256
+eval_frequency: 25 #1000 # in epochs
+train_pick_best_sample_among_k: 50
+similarity_space: latent_space # input_space, latent_space or metric_space
+
+
+diffusion_activation: identity
+num_prob_samples: 50
+diffusion_timesteps: 10
+
+diffusion_type: NonisotropicGaussianDiffusion
+diffusion_loss_type: snr #snr_triangle_inequality,mahalanobis, snr
+loss_reduction_type: l1
+if_run_as_isotropic: False
+if_sigma_n_scale: True
+diffusion_covariance_type: skeleton-diffusion # anisotropic, isotropic, skeleton-diffusion
+gamma_scheduler: cosine # mono_decrease, cosine
+beta_schedule: cosine
+sigma_n_scale: spectral
+num_iter_perepoch: null
+seed: 63485
+
+diffusion_arch: 
+    arch: Denoiser
+    use_attention: True
+    self_condition: False # True holds better results, but it takes longer to train.
+    norm_type: none
+    depth: 1
+    # resnet_block_groups = 8,
+    # learned_variance: False
+    # learned_sinusoidal_cond: False
+    # random_fourier_features: False
+    # learned_sinusoidal_dim: 16
+    # sinusoidal_pos_emb_theta: 10000
+    attn_dim_head: 32
+    attn_heads: 4
+    learn_influence: True

SkeletonDiffusion/datasets

@@ -0,0 +1 @@
+../../motion_must_go_on/datasets/

SkeletonDiffusion/environment_inference.yml

@@ -0,0 +1,19 @@
+name: skeldiff_inf
+channels:
+  - pytorch
+  - nvidia
+dependencies:
+  - python=3.10.12
+  - pytorch=2.0.1
+  - pytorch-cuda=11.8
+  - torchvision=0.15.2
+  - pyyaml=6.0.1
+  - einops=0.7.0
+  - pip
+  - pip:
+    - denoising-diffusion-pytorch==1.9.4
+  # - pyyaml=6.0.1
+  # - imageio
+  # - ipympl=0.9.3
+  # - ffmpeg
+  # - opencv

SkeletonDiffusion/inference.ipynb

@@ -0,0 +1,343 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "e927d3c2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/storage/user/yaji/yaji/NonisotropicSkeletonDiffusion/SkeletonDiffusion/src\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "os.chdir(r\"/home/stud/yaji/storage/user/yaji/NonisotropicSkeletonDiffusion/SkeletonDiffusion/src\")    \n",
+    "root_path = os.getcwd()\n",
+    "print(root_path)\n",
+    "\n",
+    "os.environ[\"CUBLAS_WORKSPACE_CONFIG\"] = \":4096:8\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "32b71ab8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/stud/yaji/miniconda3/envs/live_demo/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n",
+      "/storage/user/yaji/yaji/NonisotropicSkeletonDiffusion/SkeletonDiffusion/src/core/diffusion/base.py:184: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  @autocast(enabled = False)\n",
+      "/storage/user/yaji/yaji/NonisotropicSkeletonDiffusion/SkeletonDiffusion/src/core/diffusion/isotropic.py:72: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  @autocast(enabled = False)\n",
+      "/storage/user/yaji/yaji/NonisotropicSkeletonDiffusion/SkeletonDiffusion/src/core/diffusion/nonisotropic.py:138: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  @autocast(enabled = False)\n"
+     ]
+    }
+   ],
+   "source": [
+    "from eval_prepare_model import prepare_model, get_prediction, load_model_config_exp\n",
+    "from data import create_skeleton\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "import random\n",
+    "\n",
+    "def set_seed(seed=0):\n",
+    "    torch.use_deterministic_algorithms(True)\n",
+    "    torch.backends.cudnn.deterministic = True\n",
+    "    torch.backends.cudnn.benchmark = False\n",
+    "    np.random.seed(seed)\n",
+    "    random.seed(seed)\n",
+    "    torch.cuda.manual_seed(seed)\n",
+    "    torch.cuda.manual_seed_all(seed)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "0963a8bd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint_path = '/usr/wiss/curreli/work/my_exps/checkpoints_release/amass/diffusion/cvpr_release/checkpoints/checkpoint_150.pt'\n",
+    "# checkpoint_path = '/usr/wiss/curreli/work/my_exps/checkpoints_release/amass-mano/diffusion/cvpr_release/checkpoints/checkpoint_150.pt'\n",
+    "\n",
+    "\n",
+    "num_samples = 50"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "c6ce7b0a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "> GPU 0 ready: NVIDIA RTX A2000 12GB\n",
+      "Loading Autoencoder checkpoint: /usr/wiss/curreli/work/my_exps/checkpoints_release/amass/autoencoder/cvpr_release/checkpoints/checkpoint_300.pt ...\n",
+      "Diffusion is_ddim_sampling:  False\n",
+      "Loading Diffusion checkpoint: /usr/wiss/curreli/work/my_exps/checkpoints_release/amass/diffusion/cvpr_release/checkpoints/checkpoint_150.pt ...\n"
+     ]
+    }
+   ],
+   "source": [
+    "set_seed(seed=0)\n",
+    "\n",
+    "config, exp_folder = load_model_config_exp(checkpoint_path)\n",
+    "config['checkpoint_path'] = checkpoint_path\n",
+    "skeleton = create_skeleton(**config)   \n",
+    "\n",
+    "\n",
+    "model, device, *_ = prepare_model(config, skeleton, **config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "5c4aa1a7",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 30, 22, 3])"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# prepare input\n",
+    "# load input. It should be in meters\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "obs = np.load('/usr/wiss/curreli/work/my_exps/checkpoints_release/amass/exaple_obs.npy') # (t_past, J, 3)\n",
+    "# obs = np.load('/usr/wiss/curreli/work/my_exps/checkpoints_release/amass-mano/example_obs.npy') # (t_past, J, 3)\n",
+    "\n",
+    "\n",
+    "obs = torch.from_numpy(obs).to(device)\n",
+    "obs = obs.unsqueeze(0) # add bacth size\n",
+    "obs.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "e782b749",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1, 50, 120, 21, 3])\n",
+      "torch.Size([1, 30, 21, 3])\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 50, 120, 21, 3])"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "obs_in = skeleton.tranform_to_input_space(obs) # obs sequence contains hip joints, it has not been dropped yet. \n",
+    "pred = get_prediction(obs_in, model, num_samples=num_samples, **config) # [batch_size, n_samples, seq_length, num_joints, features]\n",
+    "print(pred.shape)\n",
+    "pred = skeleton.transform_to_metric_space(pred)\n",
+    "print(obs_in.shape)\n",
+    "pred.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d9394774",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "kpts3d = pred.cpu()[0][0]\n",
+    "import matplotlib.pyplot as plt\n",
+    "for i in range(120):\n",
+    "    plt.figure()\n",
+    "    plt.scatter(kpts3d[0,  :, 1], kpts3d[0, :, 2])\n",
+    "    plt.gca().set_aspect('equal')\n",
+    "    plt.savefig(f'../../vis/kpts3d_{i}.png')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "f91ee849",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1, 30, 21, 3])\n",
+      "torch.Size([1, 50, 120, 21, 3])\n",
+      "torch.Size([1, 30, 20, 3])\n",
+      "tensor([[0.0123, 0.0127, 0.0128, 0.0128, 0.0131, 0.0131, 0.0137, 0.0144, 0.0148,\n",
+      "         0.0149, 0.0151, 0.0153, 0.0154, 0.0158, 0.0164, 0.0168, 0.0170, 0.0171,\n",
+      "         0.0171, 0.0173, 0.0189, 0.0193, 0.0224, 0.0229, 0.0235, 0.0237, 0.0244,\n",
+      "         0.0248, 0.0261, 0.0269, 0.0270, 0.0271, 0.0274, 0.0275, 0.0280, 0.0282,\n",
+      "         0.0292, 0.0293, 0.0307, 0.0339, 0.0346, 0.0350, 0.0351, 0.0367, 0.0379,\n",
+      "         0.0386, 0.0391, 0.0395, 0.0439, 0.0523]], device='cuda:0')\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[ 4, 12, 31, 34, 27, 29,  2, 46, 20, 11, 15, 33, 21, 14,  9, 38, 41,  1,\n",
+       "         22, 35, 19, 43, 16, 48,  5, 47, 25, 40,  8, 28, 39, 45, 17, 23, 37, 18,\n",
+       "          6, 42, 49, 26, 24, 13, 36,  3, 44,  0,  7, 10, 32, 30]],\n",
+       "       device='cuda:0')"
+      ]
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# rank predictions according to Limb Stretching. We will visualize first the prediction that have lower limb stretching --> more realistic\n",
+    "from metrics.body_realism import limb_stretching_normed_mean, limb_stretching_normed_rmse\n",
+    "print(obs_in.shape)\n",
+    "print(pred.shape)\n",
+    "print(obs_in[..., 1:, :].shape)\n",
+    "# limbstretching = limb_stretching_normed_mean(pred, target=obs[..., 1:, :].unsqueeze(1), limbseq=skeleton.get_limbseq(), reduction='persample', obs_as_target=True)\n",
+    "limbstretching = limb_stretching_normed_rmse(pred, target=obs[..., 1:, :].unsqueeze(1), limbseq=skeleton.get_limbseq(), reduction='persample', obs_as_target=True)\n",
+    "limbstretching_sorted, indices =  torch.sort(limbstretching.squeeze(1), dim=-1, descending=False) \n",
+    "\n",
+    "print(limbstretching_sorted)\n",
+    "indices\n",
+    "\n",
+    "# TO DO: index predictions with these indices.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "1a9a941a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Observation shape: torch.Size([1, 30, 21, 3])\n",
+      "Prediction shape: torch.Size([1, 50, 120, 22, 3])\n",
+      "tensor([[0.0446, 0.0711, 0.0517, 0.1149, 0.0689, 0.0238, 0.0374, 0.0329, 0.0411,\n",
+      "         0.0573, 0.0565, 0.0855, 0.0375, 0.1141, 0.0402, 0.0385, 0.0564, 0.0727,\n",
+      "         0.0904, 0.0620, 0.0374, 0.0363, 0.0443, 0.0386, 0.0702, 0.0413, 0.0455,\n",
+      "         0.0468, 0.1038, 0.0691, 0.0630, 0.0320, 0.0489, 0.0422, 0.0520, 0.0756,\n",
+      "         0.0444, 0.0414, 0.0852, 0.0673, 0.0391, 0.0500, 0.0484, 0.0457, 0.0556,\n",
+      "         0.0393, 0.0674, 0.0349, 0.0392, 0.0459]], device='cuda:0')\n"
+     ]
+    }
+   ],
+   "source": [
+    "# read pred and obs from predictions/joints3d.npy\n",
+    "frames_for_half_second = 30\n",
+    "# Load the joints3d data from the saved numpy file\n",
+    "joints3d = np.load('/home/stud/yaji/storage/user/yaji/NonisotropicSkeletonDiffusion/predictions/joints3d.npy')\n",
+    "\n",
+    "# Split the data into observation and prediction parts\n",
+    "# The first frames_for_half_second frames are observations\n",
+    "obs = joints3d[:, 0, :frames_for_half_second, :, :]  # [1, num_samples, frames_for_half_second, 22, 3]\n",
+    "pred = joints3d[:, :, frames_for_half_second:, :, :]  # [1, num_samples, pred_length, 22, 3]\n",
+    "\n",
+    "# Convert to torch tensors and move to device\n",
+    "obs = torch.from_numpy(obs).to(device)\n",
+    "pred = torch.from_numpy(pred).to(device)\n",
+    "\n",
+    "print(\"Observation shape:\", obs[0, ..., 1:, :].unsqueeze(0).shape)\n",
+    "print(\"Prediction shape:\", pred.shape)\n",
+    "\n",
+    "# calculate \n",
+    "limbstretching = limb_stretching_normed_rmse(pred[..., 1:, :], target=obs[0, ..., 1:, :].unsqueeze(0), limbseq=skeleton.get_limbseq(), reduction='persample', obs_as_target=True)\n",
+    "limbstretching_sorted, indices =  torch.sort(limbstretching.squeeze(1), dim=-1, descending=False) \n",
+    "print(limbstretching)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "e51ccbee",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1, 120, 21, 3])\n",
+      "torch.Size([50, 120, 21, 3])\n",
+      "torch.Size([120, 21, 3])\n",
+      "torch.Size([10, 120, 21, 3])\n",
+      "[8, 18, 36, 30, 37, 43, 26, 17, 41, 6]\n"
+     ]
+    }
+   ],
+   "source": [
+    "from metrics.ranking import get_closest_and_nfurthest_maxapd\n",
+    "# If you see problems with the visualizations, you can remove predictions that have limb stretching > 0.04\n",
+    "# limbstretching = limb_stretching_normed_mean(pred, target=obs[..., 1:, :], limbseq=skeleton.get_limbseq(), reduction='persample', obs_as_target=True)\n",
+    "# remove batch dimension\n",
+    "y_pred = pred.squeeze(0) # [n_samples, seq_length, num_joints, features]\n",
+    "#if GT is not there, we use the first sample as GT reference i.e. the most likely closest to GT\n",
+    "y_gt = y_pred[0].unsqueeze(0) # [seq_length, num_joints, features]\n",
+    "print(y_gt.shape)\n",
+    "print(y_pred.shape)\n",
+    "pred_closest, sorted_preds, sorted_preds_idxs = get_closest_and_nfurthest_maxapd(y_pred, y_gt, nsamples=10)\n",
+    "\n",
+    "print(pred_closest.shape)\n",
+    "print(sorted_preds.shape)\n",
+    "print(sorted_preds_idxs)\n",
+    "\n",
+    "\n",
+    "\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "live_demo",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

SkeletonDiffusion/inference_filtered.ipynb

@@ -0,0 +1 @@
+ 

SkeletonDiffusion/setup.py

@@ -0,0 +1,13 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="SkeletonDiffusion",
+    version="0.1.0",
+    package_dir={"": "src"},
+    packages=find_packages(where="src"),
+    install_requires=[
+        "torch",
+        "numpy",
+    ],
+    python_requires=">=3.8",
+)

SkeletonDiffusion/src/__init__.py

@@ -0,0 +1,7 @@
+import os
+import sys
+
+# Add the src directory to the Python path
+src_path = os.path.dirname(os.path.abspath(__file__))
+if src_path not in sys.path:
+    sys.path.insert(0, src_path)

SkeletonDiffusion/src/config_utils.py

@@ -0,0 +1,62 @@
+import torch
+from torch.utils.data import DataLoader
+from ignite.handlers import Checkpoint
+import os
+
+from SkeletonDiffusion.src.utils.config import init_obj
+from SkeletonDiffusion.src.utils.reproducibility import seed_worker, seed_eval_worker, RandomStateDict
+import SkeletonDiffusion.src.data.loaders as dataset_type
+
+from SkeletonDiffusion.src.utils.load import get_latest_model_path
+from SkeletonDiffusion.src.inference_utils import create_model
+
+
+
+def create_train_dataloaders(batch_size, batch_size_eval, num_workers, skeleton, if_run_validation=True, if_resume_training=False, **config):
+    
+    random_state_manager = RandomStateDict()
+    if if_resume_training:
+        checkpoint = torch.load(config['load_path'])
+        Checkpoint.load_objects(to_load={"random_states": random_state_manager}, checkpoint=checkpoint)  
+    dataset_train = init_obj(config, 'dataset_type', dataset_type, split="train", skeleton=skeleton, **(config['data_loader_train']))
+    data_loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, worker_init_fn = seed_worker, pin_memory= True,
+                                num_workers=num_workers, generator=random_state_manager.generator)  
+     
+    if if_run_validation:
+        dataset_eval = init_obj(config, 'dataset_type', dataset_type, split="valid", skeleton=skeleton, **(config['data_loader_valid']))
+        dataset_eval_train = init_obj(config, 'dataset_type', dataset_type, split="train", skeleton=skeleton, **(config['data_loader_train_eval']))
+        data_loader_eval = DataLoader(dataset_eval, shuffle=False, worker_init_fn=seed_eval_worker, batch_size=batch_size_eval, num_workers=1, pin_memory= True)    
+        data_loader_train_eval = DataLoader(dataset_eval_train, shuffle=False, worker_init_fn=seed_eval_worker, batch_size=batch_size_eval, num_workers=1, pin_memory= True)
+    else: 
+        data_loader_eval = None
+        data_loader_train_eval = None
+    return data_loader_train, data_loader_eval, data_loader_train_eval, random_state_manager
+
+
+def flat_hydra_config(cfg):
+    """
+    Flatten the main dict categories of the Hydra config object into a single one.
+    """
+    for subconf in ['model', 'task', 'dataset', 'autoenc_arch', 'cov_matrix']:
+        if subconf in cfg:
+            cfg = {**cfg, **cfg[subconf]}
+            cfg.pop(subconf)
+    return cfg
+
+def resume_training(cfg):
+    #output folder has been already created.
+    assert 'output_log_path' in cfg
+
+    # decide whether to start from scratch (default if no checkpoints), from latest save (if checkpoints exists), or from given path (if load_path is given)
+
+    assert len(os.listdir(os.path.join(cfg['output_log_path'], 'checkpoints'))) != 0, "Checkpoints folder is empty. Please provide a valid path to load from."
+        
+    if len(cfg['load_path']) == 0:
+        #  load latest model
+        cfg['load_path'] = get_latest_model_path(os.path.join(cfg['output_log_path'], 'checkpoints'))
+        print("Loading latest epoch: ", cfg['load_path'].split('/')[-1])
+    else: 
+        output_path = os.path.dirname(os.path.dirname(cfg['load_path']))
+        assert cfg['output_log_path'] ==  output_path
+        cfg['output_log_path'] = output_path
+    return cfg

SkeletonDiffusion/src/core/__init__.py

@@ -0,0 +1,8 @@
+# Import network first
+from .network import AutoEncoder, Denoiser
+
+# Then import diffusion manager
+from .diffusion_manager import DiffusionManager
+
+# Export all
+__all__ = ['AutoEncoder', 'Denoiser', 'DiffusionManager']

SkeletonDiffusion/src/core/diffusion/__init__.py

@@ -0,0 +1,3 @@
+from SkeletonDiffusion.src.core.diffusion.isotropic import IsotropicGaussianDiffusion
+from SkeletonDiffusion.src.core.diffusion.nonisotropic import NonisotropicGaussianDiffusion
+from SkeletonDiffusion.src.core.diffusion.utils import get_cov_from_corr

SkeletonDiffusion/src/core/diffusion/base.py

@@ -0,0 +1,445 @@
+import math
+from random import random
+from functools import partial
+from collections import namedtuple
+from typing import Tuple, Optional, List, Union, Dict
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.cuda.amp import autocast
+
+from einops import reduce
+
+
+from tqdm.auto import tqdm
+
+
+# constants
+ModelPrediction =  namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
+
+# helper functions
+
+def identity(t, *args, **kwargs):
+    return t
+
+def exists(x):
+    return x is not None
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if callable(d) else d
+
+def extract(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+def linear_beta_schedule(timesteps):
+    scale = 1000 / timesteps
+    beta_start = scale * 0.0001
+    beta_end = scale * 0.02
+    return torch.linspace(beta_start, beta_end, timesteps, dtype = torch.float64)
+
+def cosine_beta_schedule(timesteps, s = 0.008):
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+    """
+    steps = timesteps + 1
+    x = torch.linspace(0, timesteps, steps, dtype = torch.float64)
+    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    return torch.clip(betas, 0, 0.999)
+
+def exp_beta_schedule(timesteps, factor=3.0):
+    steps = timesteps + 1
+    x = torch.linspace(-factor, 0, steps, dtype = torch.float64)#/timesteps
+    betas = torch.exp(x)
+    return torch.clip(betas, 0, 0.999)
+
+
+class LatentDiffusion(nn.Module):
+    def __init__(self,
+        model:torch.nn.Module, latent_size=96, diffusion_timesteps=10, diffusion_objective="pred_x0", sampling_timesteps=None, diffusion_activation='identity', 
+        silent=True, diffusion_conditioning=False, diffusion_loss_type='mse',
+        objective = 'pred_noise',
+        beta_schedule = 'cosine',
+        beta_schedule_factor=3.0,
+        ddim_sampling_eta = 0.,
+        **kwargs
+    ):
+        
+        super().__init__()
+
+        
+        if diffusion_activation == "tanh":
+            self.activation = torch.nn.Tanh()
+        elif diffusion_activation == "identity":
+            self.activation = torch.nn.Identity()
+        self.silent = silent
+        self.condition = diffusion_conditioning
+        self.loss_type = diffusion_loss_type
+        
+        self.statistics_pred = None
+        self.statistics_obs = None
+
+        
+        timesteps=diffusion_timesteps
+        objective=diffusion_objective
+        
+        
+        self.model = model
+        self.channels = self.model.channels
+        self.self_condition = self.model.self_condition
+
+        self.seq_length = latent_size
+        self.objective = objective
+
+        assert objective in {'pred_noise', 'pred_x0', 'pred_v'}, 'objective must be either pred_noise (predict noise) or pred_x0 (predict image start) or pred_v (predict v [v-parameterization as defined in appendix D of progressive distillation paper, used in imagen-video successfully])'
+
+        if beta_schedule == 'linear':
+            betas = linear_beta_schedule(timesteps)
+        elif beta_schedule == 'cosine':
+            betas = cosine_beta_schedule(timesteps)
+        elif beta_schedule == 'exp':
+            betas = exp_beta_schedule(timesteps, beta_schedule_factor)
+        else:
+            raise ValueError(f'unknown beta schedule {beta_schedule}')
+
+        alphas = 1. - betas
+        alphas_cumprod = torch.cumprod(alphas, dim=0)
+        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.)
+
+        timesteps, = betas.shape
+        self.num_timesteps = int(timesteps)
+
+        # sampling related parameters
+
+        self.sampling_timesteps = default(sampling_timesteps, timesteps) # default num sampling timesteps to number of timesteps at training
+
+        assert self.sampling_timesteps <= timesteps
+        self.is_ddim_sampling = self.sampling_timesteps < timesteps
+        self.ddim_sampling_eta = ddim_sampling_eta
+
+        # helper function to register buffer from float64 to float32
+
+        register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))
+
+        register_buffer('betas', betas)
+        register_buffer('alphas_cumprod', alphas_cumprod)
+        register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)
+        register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))       
+
+        print("Diffusion is_ddim_sampling: ", self.is_ddim_sampling)
+
+        
+   
+    def set_normalization_statistics(self, statistics_pred, statistics_obs):
+        self.statistics_pred = statistics_pred
+        self.statistics_obs = statistics_obs
+        print("Setting normalization statistics for diffusion")    
+        
+    def get_white_noise(self, x, *args, **kwargs):
+        return self.get_noise(x, *args, **kwargs)
+    
+    def get_start_noise(self, x, *args, **kwargs):
+        return self.get_white_noise(x, *args, **kwargs)
+    
+    def get_noise(self, x, *args, **kwargs):
+        """
+        x is either tensor or shape 
+        """
+        if torch.is_tensor(x):
+            return torch.randn_like(x)
+        elif isinstance(x, tuple):
+            return torch.randn(*x, *args, **kwargs) 
+        
+    #######################################################################
+    # TO SUBCLASS
+    #######################################################################
+
+    def predict_start_from_noise(self, x_t, t, noise):
+        assert 0, "Not implemented"
+        ...
+        return x_t
+
+    def predict_noise_from_start(self, x_t, t, x0):
+        assert 0, "Not implemented"
+        ...
+        return x_t
+
+    def predict_v(self, x_start, t, noise):
+        assert 0, "Not implemented"
+        ...
+        return x_start
+
+    def predict_start_from_v(self, x_t, t, v):
+        assert 0, "Not implemented"
+        ...
+        return x_t
+
+    @autocast(enabled = False)
+    def q_sample(self, x_start, t, noise=None):
+        assert 0, "Not implemented"
+        ...
+        return x_start
+
+    def q_posterior(self, x_start, x_t, t):
+        assert 0, "Not implemented"
+        ...
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_combine_mean_var_noise(self, model_mean, model_log_variance, noise):
+        assert 0, "Not implemented"
+        ...
+        return model_mean 
+    
+    def p_interpolate_mean_var_noise(self, model_mean, model_log_variance, noise, node_idx:Optional[int] = None, interpolate_factor=0.0, noise2interpolate=None, **kwargs):
+        assert 0, "Not implemented"
+        ...
+        return model_mean 
+    
+    def loss_funct(self, model_out, target, *args, **kwargs):
+        if self.loss_type == "mse":
+            loss = F.mse_loss(model_out, target, reduction = 'none')
+        elif self.loss_type == 'l1':
+            loss = F.l1_loss(model_out, target, reduction = 'none')
+        else: 
+            assert 0, "Not implemented"
+        return loss
+    
+    ########################################################################
+    # NETWORK INTERFACE
+    ######################################################################### 
+
+
+    def model_predictions(self, x, t, x_self_cond = None, x_cond=None, clip_x_start = False, rederive_pred_noise = False):
+        model_output = self.feed_model(x, t, x_self_cond=x_self_cond, x_cond=x_cond)
+        maybe_clip = partial(torch.clamp, min = -1., max = 1.) if clip_x_start else identity
+
+        if self.objective == 'pred_noise':
+            pred_noise = model_output
+            x_start = self.predict_start_from_noise(x, t, pred_noise)
+            x_start = maybe_clip(x_start)
+
+            if clip_x_start and rederive_pred_noise:
+                pred_noise = self.predict_noise_from_start(x, t, x_start)
+
+        elif self.objective == 'pred_x0':
+            x_start = model_output
+            x_start = maybe_clip(x_start)
+            pred_noise = self.predict_noise_from_start(x, t, x_start)
+
+        elif self.objective == 'pred_v':
+            v = model_output
+            x_start = self.predict_start_from_v(x, t, v)
+            x_start = maybe_clip(x_start)
+            pred_noise = self.predict_noise_from_start(x, t, x_start)
+        return ModelPrediction(pred_noise, x_start) 
+    
+    def feed_model(self, x, t, x_self_cond = None, x_cond=None): 
+        if self.condition:
+            assert x_cond is not None
+            if x.shape[0] > x_cond.shape[0]:
+                # training with multiple samples
+                x_cond = x_cond.repeat_interleave( int(x.shape[0]/x_cond.shape[0]), 0)
+            model_in = x 
+        else: 
+            model_in = x
+
+        model_output = self.model(model_in, t, x_self_cond, x_cond)
+        model_output = self.activation(model_output)
+        return model_output
+    
+    ########################################################################
+    # FORWARD PROCESS
+    #########################################################################  
+    
+    
+    def p_losses(self, x_start, t, noise = None, x_cond=None, n_train_samples=1):
+        b, c, n = x_start.shape
+        if n_train_samples > 1:
+            x_start = x_start.repeat_interleave(n_train_samples, dim=0)
+            t = t.repeat_interleave(n_train_samples, dim=0)
+            if x_cond is not None:
+                x_cond = x_cond.repeat_interleave(n_train_samples, dim=0)
+        noise = default(noise, self.get_white_noise(x_start, t)) # noise for timesteps t
+
+        # noise sample
+
+        x = self.q_sample(x_start = x_start, t = t, noise = noise)
+
+        # if doing self-conditioning, 50% of the time, predict x_start from current set of times
+        # and condition with unet with that
+        # this technique will slow down training by 25%, but seems to lower FID significantly
+
+        x_self_cond = None
+        if self.self_condition and random.random() < 0.5:
+            with torch.no_grad():
+                x_self_cond = self.model_predictions(x, t, x_cond=x_cond).pred_x_start
+                x_self_cond.detach_()
+
+        # predict and take gradient step
+        model_out = self.feed_model(x, t, x_self_cond=x_self_cond, x_cond=x_cond)
+
+        if self.objective == 'pred_noise':
+            target = noise
+        elif self.objective == 'pred_x0':
+            target = x_start
+        elif self.objective == 'pred_v':
+            v = self.predict_v(x_start, t, noise)
+            target = v
+        else:
+            raise ValueError(f'unknown objective {self.objective}')
+        loss = self.loss_funct(model_out, target, t)
+        loss = reduce(loss, 'b ... -> b', 'mean') # [batch*n_train_samples, Nodes, latent_dim] -> [batch*n_train_samples]
+
+        return loss,  extract(self.loss_weight, t.view(b, -1)[:, 0], loss.shape[0:1]), model_out
+    
+    def forward(self, x, *args, x_cond=None, **kwargs):
+        b, c, n, device, seq_length, = *x.shape, x.device, self.seq_length
+        assert n == seq_length, f'seq length must be {seq_length}'
+        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
+
+        return self.p_losses(x, t, *args, x_cond=x_cond, **kwargs)
+    
+    
+    ########################################################################
+    # REVERSE PROCESS
+    #########################################################################
+    
+    def p_mean_variance(self, x, t, x_self_cond = None, x_cond=None, clip_denoised = True):
+        preds = self.model_predictions(x, t, x_self_cond, x_cond=x_cond)
+        x_start = preds.pred_x_start
+
+        if clip_denoised:
+            x_start.clamp_(-1., 1.)
+
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start = x_start, x_t = x, t = t)
+        return model_mean, posterior_variance, posterior_log_variance, x_start
+
+    @torch.no_grad()
+    def p_sample(self, x, t: int, x_self_cond = None, clip_denoised = True, sampling_noise=None, *args, if_interpolate=False, noise2interpolate=None, interpolation_kwargs:Dict=None, **kwargs):
+        b, *_, device = *x.shape, x.device
+        batched_times = torch.full((b,), t, device = x.device, dtype = torch.long)
+        model_mean, _, model_log_variance, x_start = self.p_mean_variance(x = x, t = batched_times, x_self_cond = x_self_cond, clip_denoised = clip_denoised, *args, **kwargs)
+        
+        if sampling_noise is not None and t > 0:
+            noise = sampling_noise[:, sampling_noise.shape[1]-t]
+        else:
+            noise = self.get_white_noise(x) if t > 0 else 0. # no noise if t == 0
+        
+        if if_interpolate and t > 0:
+            noise2 = noise2interpolate[:, sampling_noise.shape[1]-t]
+            assert noise2.shape == noise.shape
+            pred_img = self.p_interpolate_mean_var_noise(model_mean, model_log_variance, noise, noise2, **interpolation_kwargs)
+        else:
+            pred_img = self.p_combine_mean_var_noise(model_mean, model_log_variance, noise)
+        return pred_img, x_start, noise, model_mean
+
+    @torch.no_grad()
+    def p_sample_loop(self, shape, x_cond=None, start_noise=None, sampling_noise=None, return_sampling_noise=False, return_timages=False, **kwargs):
+        batch, device = shape[0], self.betas.device
+        if start_noise is not None:
+            assert start_noise.shape == shape, f"Shape mismatch: {start_noise.shape} != {shape}"
+            img = start_noise
+            noise = start_noise.clone()
+        else:
+            img = self.get_start_noise(shape, device=device)
+            noise = img.clone()
+            
+        if sampling_noise is not None:
+            assert sampling_noise.shape[2:] == shape[1:], f"Shape mismatch: {start_noise.shape} != {shape}"
+            assert sampling_noise.shape[0] == shape[0], f"Shape mismatch: {start_noise.shape} != {shape}"
+            assert sampling_noise.shape[1] == self.num_timesteps - 1
+
+        x_start = None
+        imgs = []
+        noise_t = []
+        mean_t = []
+        if not self.silent:
+            print(f"Evaluation with {len(range(0, self.num_timesteps))} diffusion steps")
+        for t in reversed(range(0, self.num_timesteps)): #, desc = 'sampling loop time step', total = self.num_timesteps):
+            self_cond = x_start if self.self_condition else None
+            img, x_start, nt, model_mean = self.p_sample(img, t, self_cond, x_cond=x_cond, sampling_noise=sampling_noise, **kwargs)
+            if return_sampling_noise and t!=0:
+                noise_t.append(nt)
+                mean_t.append(model_mean)
+            if return_timages and t!=0:
+                imgs.append(img)
+
+
+        if return_sampling_noise:
+            noise_t = torch.stack(noise_t, dim=1)
+            mean_t = torch.stack(mean_t, dim=1)
+        if return_timages:
+            print("Returning timages")
+            imgs = torch.stack(imgs, dim=1)
+
+        if return_sampling_noise:
+            if return_timages:
+                noise = (noise, noise_t, imgs)
+            else:
+                noise = (noise, noise_t, mean_t)
+        else:
+            if return_timages:
+                noise = (noise, imgs)
+        return img, noise
+   
+    @torch.no_grad()
+    def ddim_sample(self, shape, clip_denoised = True, x_cond=None, start_noise=None): 
+        batch, device, total_timesteps, sampling_timesteps, eta, objective = shape[0], self.betas.device, self.num_timesteps, self.sampling_timesteps, self.ddim_sampling_eta, self.objective
+
+        times = list(reversed(times.int().tolist()))
+        time_pairs = list(zip(times[:-1], times[1:])) # [(T-1, T-2), (T-2, T-3), ..., (1, 0), (0, -1)]
+        
+        if start_noise is not None:
+            assert start_noise.shape == shape
+            img = start_noise
+            noise = start_noise.clone()
+        else:
+            img = torch.randn(shape, device=device)
+            noise = img.clone()
+
+        imgs = []
+
+        x_start = None
+        if not self.silent:
+            print(f"Evaluation with {len(time_pairs)} diffusion steps")
+        for time, time_next in tqdm(time_pairs, desc = 'sampling loop time step'):
+            time_cond = torch.full((batch,), time, device=device, dtype=torch.long)
+            self_cond = x_start if self.self_condition else None
+            pred_noise, x_start, *_ = self.model_predictions(img, time_cond, self_cond, x_cond=x_cond, clip_x_start = clip_denoised)
+
+            if time_next < 0:
+                img = x_start
+                # imgs.append(img)
+
+            alpha = self.alphas_cumprod[time]
+            alpha_next = self.alphas_cumprod[time_next]
+            sqrt_alpha_next = self.sqrt_alphas_cumprod[time_next]
+
+            sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt() #eta*beta_t_tilde
+            c = (1 - alpha_next - sigma ** 2).sqrt()
+
+            noise = torch.randn_like(img)
+
+            img = x_start * sqrt_alpha_next + \
+                  c * pred_noise + \
+                  sigma * noise
+            # imgs.append(img)
+
+        # imgs = torch.stack(imgs, dim=1)
+        return img, noise
+    
+    
+    @torch.no_grad()
+    def sample(self, batch_size = 16, *args, **kwargs):
+        seq_length, channels = self.seq_length, self.channels
+        sample_fn = self.p_sample_loop if not self.is_ddim_sampling else self.ddim_sample
+        return sample_fn((batch_size, channels, seq_length),*args, **kwargs)
+
+

SkeletonDiffusion/src/core/diffusion/isotropic.py

@@ -0,0 +1,104 @@
+import torch
+from torch.cuda.amp import autocast
+
+from SkeletonDiffusion.src.core.diffusion.base import LatentDiffusion, extract, default
+
+class IsotropicGaussianDiffusion(LatentDiffusion):
+    def __init__(self, **kwargs):
+        super().__init__( **kwargs)
+        register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+
+        register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - self.alphas_cumprod))
+        register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - self.alphas_cumprod))
+        register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / self.alphas_cumprod))
+        register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / self.alphas_cumprod - 1))
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+
+        posterior_variance = self.betas * (1. - self.alphas_cumprod_prev) / (1. - self.alphas_cumprod)
+        alphas = 1. - self.betas
+
+        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
+
+        register_buffer('posterior_variance', posterior_variance)
+
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+
+        register_buffer('posterior_log_variance_clipped', torch.log(posterior_variance.clamp(min =1e-20)))
+        register_buffer('posterior_mean_coef1', self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1. - self.alphas_cumprod))
+        register_buffer('posterior_mean_coef2', (1. - self.alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - self.alphas_cumprod))
+        
+        # calculate loss weight
+        snr = self.alphas_cumprod / (1 - self.alphas_cumprod)
+
+        if self.objective == 'pred_noise':
+            loss_weight = torch.ones_like(snr)
+        elif self.objective == 'pred_x0':
+            loss_weight = snr
+        elif self.objective == 'pred_v':
+            loss_weight = snr / (snr + 1)
+
+        register_buffer('loss_weight', loss_weight)
+
+    ########################################################################
+    # FORWARD PROCESS
+    #########################################################################  
+
+    def predict_start_from_noise(self, x_t, t, noise):
+        return (
+            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
+            extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
+        )
+
+    def predict_noise_from_start(self, x_t, t, x0):
+        return (
+            (extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0) / \
+            extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+        )
+
+    def predict_v(self, x_start, t, noise):
+        return (
+            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * noise -
+            extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * x_start
+        )
+
+    def predict_start_from_v(self, x_t, t, v):
+        return (
+            extract(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t -
+            extract(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v
+        )
+        
+    @autocast(enabled = False)
+    def q_sample(self, x_start, t, noise=None):
+        noise = default(noise, lambda: self.get_white_noise(x_start))
+
+        return (
+            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+            extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+        
+    ########################################################################
+    # REVERSE PROCESS
+    #########################################################################
+        
+    def q_posterior(self, x_start, x_t, t):
+        posterior_mean = (
+            extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
+            extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_combine_mean_var_noise(self, model_mean, model_log_variance, noise):
+        return model_mean + (0.5 * model_log_variance).exp() * noise 
+    
+    def interpolate_noise(self, noise1, noise2, interpolate_funct=None, **kwargs):
+        interpolated_noise = interpolate_funct(noise1, noise2)
+        return interpolated_noise
+        
+    def p_interpolate_mean_var_noise(self, model_mean, model_log_variance, noise, noise2interpolate=None, **kwargs):
+        interpolated_noise = self.interpolate_noise(noise, noise2interpolate,**kwargs) 
+        return model_mean + (0.5 * model_log_variance).exp() * interpolated_noise 
+    

SkeletonDiffusion/src/core/diffusion/nonisotropic.py

@@ -0,0 +1,213 @@
+import torch
+from torch.cuda.amp import autocast
+from .base import LatentDiffusion, extract, default
+
+def extract_matrix(matrix, t, x_shape):
+    b, *_ = t.shape
+    T, N, *_ = matrix.shape
+    out = torch.index_select(matrix, 0, t)
+    out = out.reshape(b, *out.shape[1:])
+    while len(x_shape) > len(out.shape):
+        out = out.unsqueeze(-1)
+    return out
+
+def verify_noise_scale(diffusion):
+    N, *_ = diffusion.Lambda_N.shape
+    alphas = 1 - diffusion.betas
+    noise = diffusion.get_noise((2000, diffusion.num_timesteps, N))
+    zeta_noise = torch.sqrt(diffusion.Lambda_t.unsqueeze(0)) * noise
+    print("current: ", (zeta_noise**2).sum(-1).mean(0))
+    print("original standard gaussian diffusion: ",(1-alphas) * zeta_noise.shape[-1])
+
+def compute_covariance_matrices(diffusion: torch.nn.Module,  Lambda_N: torch.Tensor,  diffusion_covariance_type='ani-isotropic', gamma_scheduler = 'cosine'):
+    N, *_  = Lambda_N.shape
+    alphas = 1. - diffusion.betas
+    def _alpha_sumprod(alphas, t):
+        return torch.sum(torch.cumprod(torch.flip(alphas[:t+1], [0]), dim=0))
+    alphas_sumprod = torch.stack([_alpha_sumprod(alphas, t) for t in range(len(alphas))], dim=0)
+    diffusion.alphas_sumprod = alphas_sumprod
+    if diffusion_covariance_type == 'isotropic':
+        assert (Lambda_N == 0).all()
+        Lambda_t = (1-alphas).unsqueeze(-1) # (Tdiff, N)
+        Lambda_bar_t = (1-diffusion.alphas_cumprod.unsqueeze(-1))
+        Lambda_bar_t_prev = torch.cat([torch.zeros(1).unsqueeze(0), Lambda_bar_t[:-1]], dim=0)
+    elif diffusion_covariance_type == 'anisotropic': 
+        Lambda_t = (1-alphas.unsqueeze(-1))*Lambda_N # (Tdiff, N)
+        Lambda_bar_t = (1-diffusion.alphas_cumprod.unsqueeze(-1))*Lambda_N
+        Lambda_bar_t_prev = (1-diffusion.alphas_cumprod_prev.unsqueeze(-1))*Lambda_N
+    elif diffusion_covariance_type == 'skeleton-diffusion':
+        if gamma_scheduler== 'cosine':
+            gammas = 1 - alphas
+        elif gamma_scheduler == 'mono_decrease':
+            gammas = 1 - torch.arange(0, diffusion.num_timesteps)/diffusion.num_timesteps
+        else: 
+            assert 0, "Not implemented"
+        Lambda_I = Lambda_N - 1
+        gammas_bar = (1-alphas)*gammas
+        gammas_tilde = diffusion.alphas_cumprod*torch.cumsum(gammas_bar/diffusion.alphas_cumprod, dim=-1)
+        Lambda_t = Lambda_I.unsqueeze(0)*gammas_bar.unsqueeze(-1) + (1-alphas).unsqueeze(-1) # (Tdiff, N)
+        Lambda_bar_t = Lambda_I.unsqueeze(0)*gammas_tilde.unsqueeze(-1)  + (1-diffusion.alphas_cumprod.unsqueeze(-1))
+        Lambda_bar_t_prev = torch.cat([torch.zeros(N).unsqueeze(0), Lambda_bar_t[:-1]], dim=0) # we start from det so it must be zero for t=-1
+    else:
+        assert 0, "Not implemented"
+        
+    return Lambda_t, Lambda_bar_t, Lambda_bar_t_prev
+
+
+class NonisotropicGaussianDiffusion(LatentDiffusion):
+    def __init__(self, Sigma_N: torch.Tensor, Lambda_N: torch.Tensor, U: torch.Tensor, diffusion_covariance_type='skeleton-diffusion', loss_reduction_type='l1', gamma_scheduler = 'cosine',  **kwargs):
+        super().__init__( **kwargs)
+        alphas = 1. - self.betas
+
+
+        N, _ = Sigma_N.shape
+        register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))
+        register_buffer('Lambda_N', Lambda_N)
+        register_buffer('Sigma_N', Sigma_N)
+        self.set_rotation_matrix(U)
+
+
+        Lambda_t, Lambda_bar_t, Lambda_bar_t_prev = compute_covariance_matrices(diffusion=self, Lambda_N=Lambda_N,  
+                                                                                diffusion_covariance_type=diffusion_covariance_type, gamma_scheduler=gamma_scheduler)
+
+        def create_diagonal_matrix(diagonal_vector):
+          return torch.stack([torch.diag(diag)  for diag in diagonal_vector], dim=0) # [T, N, N]
+        ######### forward , for training and inference #####################
+        #predict_noise_from_start
+        inv_sqrt_Lambda_bar = 1/torch.sqrt(Lambda_bar_t)
+        inv_sqrt_Lambda_bar_sqrt_alphas_cumprod = (1/torch.sqrt(Lambda_bar_t))*self.sqrt_alphas_cumprod.unsqueeze(-1)
+        register_buffer('inv_sqrt_Lambda_bar_mmUt', create_diagonal_matrix(inv_sqrt_Lambda_bar)@self.U_transposed.unsqueeze(0))
+        register_buffer('inv_sqrt_Lambda_bar_sqrt_alphas_cumprod_mmUt', create_diagonal_matrix(inv_sqrt_Lambda_bar_sqrt_alphas_cumprod)@self.U_transposed.unsqueeze(0))
+        #predict_start_from_noise
+        sqrt_Lambda_bar = torch.sqrt(Lambda_bar_t)
+        sqrt_Lambda_bar_sqrt_recip_alphas_cumprod = torch.sqrt(Lambda_bar_t/self.alphas_cumprod.unsqueeze(-1))
+        register_buffer('Umm_sqrt_Lambda_bar_t', U.unsqueeze(0)@create_diagonal_matrix(sqrt_Lambda_bar))
+        register_buffer('Umm_sqrt_Lambda_bar_t_sqrt_recip_alphas_cumprod', U.unsqueeze(0)@create_diagonal_matrix(sqrt_Lambda_bar_sqrt_recip_alphas_cumprod))
+
+        ######### q_posterior , for reverse process #####################
+        #q_posterior
+        Lambda_posterior_t = Lambda_t*Lambda_bar_t_prev*(1/Lambda_bar_t)
+        sqrt_alphas_cumprod_prev = torch.sqrt(self.alphas_cumprod_prev)
+        register_buffer('Lambda_posterior', Lambda_posterior_t)
+        register_buffer('Lambda_posterior_log_variance_clipped', torch.log(Lambda_posterior_t.clamp(min =1e-20)))
+        
+        posterior_mean_coef1_x0 = sqrt_alphas_cumprod_prev.unsqueeze(-1).unsqueeze(-1)*(U.unsqueeze(0)@create_diagonal_matrix((1/Lambda_bar_t)*Lambda_t)@self.U_transposed.unsqueeze(0))
+        posterior_mean_coef2_xt = torch.sqrt(alphas).unsqueeze(-1).unsqueeze(-1)*(U.unsqueeze(0)@create_diagonal_matrix((1/Lambda_bar_t)*Lambda_bar_t_prev)@self.U_transposed.unsqueeze(0))
+        register_buffer('posterior_mean_coef1_x0', posterior_mean_coef1_x0)
+        register_buffer('posterior_mean_coef2_xt', posterior_mean_coef2_xt)
+
+        ######### loss #####################
+        self.loss_reduction_type = loss_reduction_type         
+        sqrt_recip_Lambda_bar_t = torch.sqrt(1. / Lambda_bar_t)           
+        register_buffer('mahalanobis_S_sqrt_recip', create_diagonal_matrix(sqrt_recip_Lambda_bar_t)@self.U_transposed.unsqueeze(0))
+
+        if self.objective == 'pred_noise':
+            loss_weight = torch.ones_like(alphas)
+        elif self.objective == 'pred_x0':
+            loss_weight = self.alphas_cumprod
+        elif self.objective == 'pred_v':
+            assert 0, "Not implemented"
+            # loss_weight = snr / (snr + 1)
+        register_buffer('loss_weight', loss_weight)
+    
+        assert not  len(self.mahalanobis_S_sqrt_recip.shape) == 1
+            
+
+    ########################################################################
+    # CLASS FUNCTIONS
+    #########################################################################  
+    
+    def set_rotation_matrix(self, U:torch.Tensor):
+        register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))
+        register_buffer('U', U)    
+        register_buffer('U_transposed', U.t())    
+        
+    def check_eigh(self):
+        return torch.isclose(self.U@torch.diag(self.Lambda_N)@self.U_transposed,self.Sigma_N)#.all(), "U@Lambda_N@U^T must be equal to Sigma_N"
+
+    def get_anisotropic_noise(self, x, *args, **kwargs):
+        """
+        x is either tensor or shape 
+        """
+        return self.get_noise(x, *args, **kwargs)*self.Lambda_N.unsqueeze(-1)
+    
+    ########################################################################
+    # FORWARD PROCESS
+    #########################################################################        
+
+    @autocast(enabled = False)
+    def q_sample(self, x_start, t, noise=None):
+        noise = default(noise, lambda: self.get_white_noise(x_start))
+
+        return (
+            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+            extract_matrix(self.Umm_sqrt_Lambda_bar_t, t, x_start.shape) @ noise
+        )
+    # for inference    
+    def predict_start_from_noise(self, x_t, t, noise):
+        return (
+            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
+            extract_matrix(self.Umm_sqrt_Lambda_bar_t_sqrt_recip_alphas_cumprod, t, x_t.shape) @ noise
+        )
+    # for inference    
+    def predict_noise_from_start(self, x_t, t, x0):
+        return (
+            extract_matrix(self.inv_sqrt_Lambda_bar_mmUt, t, x_t.shape)@x_t -\
+                extract_matrix(self.inv_sqrt_Lambda_bar_sqrt_alphas_cumprod_mmUt, t, x_t.shape)@x0
+        )
+        
+    ########################################################################
+    # LOSS
+    ######################################################################### 
+       
+    def mahalanobis_dist(self, matrix, vector):
+        return (matrix@vector).abs() # check shape
+    
+    def loss_funct(self, model_out, target, t):
+        difference = target - model_out if self.objective == 'pred_noise' else model_out - target
+
+        loss = self.mahalanobis_dist(extract_matrix(self.mahalanobis_S_sqrt_recip, t, difference.shape), difference)
+        if self.loss_reduction_type == 'l1': 
+            loss = loss
+        elif self.loss_reduction_type == 'mse': 
+            loss = loss**2
+        else:
+            assert 0, "Not implemented"
+        return loss
+
+    ########################################################################
+    # REVERSE PROCESS
+    #########################################################################
+        
+    def q_posterior_mean(self, x_start, x_t, t):
+         return (
+                extract_matrix(self.posterior_mean_coef1_x0, t, x_t.shape) @ x_start +
+                extract_matrix(self.posterior_mean_coef2_xt, t, x_t.shape) @ x_t
+            )
+    
+    def q_posterior(self, x_start, x_t, t):
+        posterior_mean = self.q_posterior_mean(x_start, x_t, t)
+        posterior_variance = extract_matrix(self.Lambda_posterior, t, x_t.shape)
+        posterior_log_variance_clipped = extract_matrix(self.Lambda_posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_combine_mean_var_noise(self, model_mean, posterior_log_variance, noise):
+        """ mean is in not diagonal coordinate system, posterior_log_variance is in diagonal coordinate system"""
+        return model_mean + self.U@((0.5 * posterior_log_variance).exp() * noise)
+    
+
+    ########################################################################
+    # INTERPOLATION
+    #########################################################################
+
+    def interpolate_noise(self, noise1, noise2, posterior_log_variance=None, interpolate_funct=None):
+        noise1 = self.U@((0.5 * posterior_log_variance).exp() * noise1)
+        noise2 = self.U@((0.5 * posterior_log_variance).exp() * noise2)
+        interpolated_noise = interpolate_funct(noise1, noise2)
+        return interpolated_noise
+        
+    
+    def p_interpolate_mean_var_noise(self, model_mean, model_log_variance, noise, noise2interpolate=None, **kwargs):
+        interpolated_noise = self.interpolate_noise(noise, noise2interpolate, posterior_log_variance=model_log_variance, **kwargs) 
+        return model_mean + interpolated_noise # (0.5 * model_log_variance).exp()
+

SkeletonDiffusion/src/core/diffusion/utils.py

@@ -0,0 +1,125 @@
+import torch
+
+def dim_null_space(matrix):
+    assert matrix.shape[-1] == matrix.shape[-2], "Matrix must be square"
+    # rank = torch.linalg.matrix_rank(matrix) This is not set to accuracy of PYTORCH float32
+    # 1.0 + eps != 1.0
+    # torhc.tensor(1.0) + 0.7e-7!= torhc.tensor(1.0)
+    return torch.sum(torch.linalg.eigh(matrix)[0].abs() < 0.7e-7)
+        
+def is_positive_def(matrix):
+    #M is symmetric or Hermitian, and all its eigenvalues are real and positive.
+    assert  torch.allclose(matrix.transpose(-1, -2), matrix), "Matrix must be symmetric"
+    eigenvalues = torch.linalg.eigvals(matrix)
+    is_pos_def = (torch.real(eigenvalues)> 0).all()
+    if is_pos_def:
+        assert torch.isreal(eigenvalues).all(), "Eigenvalues must be real"
+    return (torch.real(eigenvalues)> 0).all()
+
+def make_positive_definite(matrix, epsilon=1e-6, if_submin=False):
+
+    eigenvalues = torch.linalg.eigvals(matrix)
+    # assert torch.isreal(eigenvalues).all()
+    if is_positive_def(matrix):
+        print("Input Matrix was positive Definitive without adding spectral norm to the diagonal")
+        return matrix
+
+    eigenvalues = torch.real(eigenvalues)
+    if not if_submin:
+        max_eig = eigenvalues.abs().max() #
+        pos_def_matrix = matrix + torch.eye(matrix.shape[0])*(max_eig + epsilon)
+    else:
+        min_eig = eigenvalues.min()
+        pos_def_matrix = matrix + torch.eye(matrix.shape[0])*(- min_eig + epsilon)
+    assert  dim_null_space(pos_def_matrix) == 0
+    return pos_def_matrix
+
+def normalize_cov(Sigma_N:torch.Tensor, Lambda_N:torch.Tensor, U:torch.Tensor, if_sigma_n_scale=True, sigma_n_scale='spectral', **kwargs):
+    N, _ = Sigma_N.shape
+    assert Lambda_N.shape == (N,)
+    assert U.shape == (N, N)
+
+    if if_sigma_n_scale:
+        # decrease the scale of Sigma_N to make it more similar to the identity matrix
+        if sigma_n_scale == 'spectral':
+            relative_scale_factor = Lambda_N.max()
+        else:
+            if sigma_n_scale == 'frob':
+                relative_scale_factor = Lambda_N.sum()/N
+            else:
+                assert 0, "Not implemented"
+        
+        Lambda_N = Lambda_N/relative_scale_factor
+        
+        Sigma_N = Sigma_N/relative_scale_factor
+        cond = U @ torch.diag(Lambda_N) @ U.mT
+        assert torch.isclose(Sigma_N, cond, atol=1e-06).all(), "Sigma_N must be equal to U @ Lambda_N @ U.t()"
+        # Sigma_N[Sigma_N>0.] = (Sigma_N + Sigma_N.t())[Sigma_N>0.]/2           
+    cond = Lambda_N>0.7e-7
+    assert (cond).all(), f"Lambda_N must be positive definite: {Lambda_N}"
+    assert is_positive_def(Sigma_N), "Sigma_N must be positive definite" 
+    # print("Frobenius Norm of SigmaN: ", torch.linalg.matrix_norm(Sigma_N, ord='fro').mean().item(), "Spectral Norm of SigmaN: ", Lambda_N.max(dim=-1)[0].mean().item())
+    return Sigma_N, Lambda_N
+
+
+def get_cov_from_corr(correlation_matrix: torch.Tensor, if_sigma_n_scale=True, sigma_n_scale='spectral', if_run_as_isotropic=False, diffusion_covariance_type='skeleton-diffusion', **kwargs):
+    N, _ = correlation_matrix.shape
+    
+    if if_run_as_isotropic:
+        if diffusion_covariance_type == 'skeleton-diffusion':
+            Lambda_N = torch.ones(N, device=correlation_matrix.device)
+            Sigma_N = torch.zeros_like(correlation_matrix)
+            U = torch.eye(N, device=correlation_matrix.device)
+        elif diffusion_covariance_type == 'anisotropic':
+            Lambda_N = torch.ones(N, device=correlation_matrix.device)
+            Sigma_N = torch.eye(N, device=correlation_matrix.device)
+            U = torch.eye(N, device=correlation_matrix.device)
+        else: 
+            Lambda_N = torch.zeros(N, device=correlation_matrix.device)
+            Sigma_N = torch.zeros_like(correlation_matrix)
+            U = torch.eye(N, device=correlation_matrix.device)
+    else: 
+        Sigma_N = make_positive_definite(correlation_matrix)
+        Lambda_N, U = torch.linalg.eigh(Sigma_N, UPLO='L')
+
+        Sigma_N, Lambda_N = normalize_cov(Sigma_N=Sigma_N, Lambda_N=Lambda_N, U=U, if_sigma_n_scale=if_sigma_n_scale, sigma_n_scale=sigma_n_scale, **kwargs)
+    return Sigma_N, Lambda_N, U
+
+
+def verify_noise_scale(diffusion):
+    N, *_ = diffusion.Lambda_N.shape
+    alphas = 1 - diffusion.betas
+    noise = diffusion.get_noise((2000, diffusion.num_timesteps, N))
+    zeta_noise = torch.sqrt(diffusion.Lambda_t.unsqueeze(0)) * noise
+    print("current: ", (zeta_noise**2).sum(-1).mean(0))
+    print("original standard gaussian diffusion: ",(1-alphas) * zeta_noise.shape[-1])
+
+
+def plot_matrix(matrix):
+    import matplotlib
+    import matplotlib.pyplot as plt
+    from matplotlib.colors import ListedColormap
+    import numpy as np
+
+    Sigma_N = matrix.cpu().clone().numpy()
+    color = 'Purples'
+    cmap = matplotlib.colormaps[color].set_bad("white")
+    # colormap_r = ListedColormap(cmap.colors[::-1])
+
+    fig, ax = plt.subplots(1,1, figsize=(6, 6),sharex=True,  subplot_kw=dict(box_aspect=1),)
+    # cax = fig.add_axes([0.93, 0.15, 0.01, 0.7])  # Adjust the position and size of the colorbar
+    # for i, ax in enumerate(axes):
+    vmax = Sigma_N.max()
+    Sigma_N[Sigma_N <=0.0000] = np.nan
+    im = ax.imshow(Sigma_N, cmap=color, vmin=0., vmax=vmax)
+    # ax.set_xticks(np.arange(len(Sigma_N)))
+    # ax.set_xticklabels(labels=list(skeleton.node_dict.values()), rotation=45, ha="right",
+    #         rotation_mode="anchor")
+    # ax.set_yticks(np.arange(len(Sigma_N)))
+    # ax.set_yticklabels(labels=list(skeleton.node_dict.values()), rotation=45, ha="right",
+    #         rotation_mode="anchor")
+    # ax.set_title(list(method2sigman.keys())[i])
+    fig.colorbar(im, cmap=cmap)
+#     plt.title('Adjancecy Matrix')
+    plt.show()
+    # fig.savefig("../paper_plots/sigmaN.pdf", format="pdf", bbox_inches="tight")

SkeletonDiffusion/src/core/diffusion_manager.py

@@ -0,0 +1,45 @@
+from typing import Tuple, Optional, List, Union, Dict, Any
+import torch
+
+from SkeletonDiffusion.src.core.network.nn import Denoiser
+from SkeletonDiffusion.src.core.diffusion import IsotropicGaussianDiffusion, NonisotropicGaussianDiffusion, get_cov_from_corr
+
+
+class DiffusionManager():
+    def __init__(self, diffusion_type: str='IsotropicGaussianDiffusion', skeleton=None, covariance_matrix_type: str = 'adjacency', 
+                 reachability_matrix_degree_factor=0.5, reachability_matrix_stop_at=0, if_sigma_n_scale=True, sigma_n_scale='spectral', if_run_as_isotropic=False, 
+                 **kwargs):
+        
+        model = self.get_network(**kwargs)
+        self.diffusion_type = diffusion_type
+
+        if diffusion_type == 'NonisotropicGaussianDiffusion':
+            # define SigmaN 
+            if covariance_matrix_type == 'adjacency':
+                correlation_matrix = skeleton.adj_matrix
+            elif covariance_matrix_type == 'reachability':
+                correlation_matrix = skeleton.reachability_matrix(factor=reachability_matrix_degree_factor, stop_at=reachability_matrix_stop_at)
+            else: 
+                assert 0, "Not implemented"
+            N, *_ = correlation_matrix.shape
+
+            Sigma_N, Lambda_N, U = get_cov_from_corr(correlation_matrix=correlation_matrix,  if_sigma_n_scale=if_sigma_n_scale, sigma_n_scale=sigma_n_scale, if_run_as_isotropic=if_run_as_isotropic, **kwargs)            
+            self.diffusion = NonisotropicGaussianDiffusion(Sigma_N=Sigma_N, Lambda_N=Lambda_N, U=U, model=model, **kwargs)
+        elif diffusion_type == 'IsotropicGaussianDiffusion':
+            self.diffusion =  IsotropicGaussianDiffusion(model=model, **kwargs)
+        else:
+            assert 0, f"{diffusion_type} Not implemented"
+        
+    def get_diffusion(self):
+        return self.diffusion
+        
+    def get_network(self, num_nodes, diffusion_conditioning=False, latent_size=96, node_types: torch.Tensor = None, diffusion_arch=Dict[str, Any], **kwargs): 
+
+        if diffusion_conditioning:
+            cond_dim = latent_size
+        else: 
+            cond_dim = 0
+        
+        model = Denoiser(dim=latent_size, cond_dim=cond_dim, out_dim=latent_size, channels=num_nodes, num_nodes=num_nodes, node_types=node_types,**diffusion_arch)
+
+        return model

SkeletonDiffusion/src/core/network/__init__.py

@@ -0,0 +1,3 @@
+from .nn import AutoEncoder, Denoiser
+
+__all__ = ['AutoEncoder', 'Denoiser']

SkeletonDiffusion/src/core/network/layers/__init__.py

@@ -0,0 +1,3 @@
+from .graph_structural import StaticGraphLinear
+from .recurrent import StaticGraphGRU, GraphGRUState, StaticGraphLSTM, GraphLSTMState
+from .attention import Attention, ResnetBlock, Residual, PreNorm, RMSNorm

SkeletonDiffusion/src/core/network/layers/attention.py

@@ -0,0 +1,138 @@
+import torch
+from torch import nn, einsum, Tensor
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn, einsum
+
+from .graph_structural import StaticGraphLinear
+
+
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, *args, **kwargs):
+        return self.fn(x, *args, **kwargs) + x
+    
+class LayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.norm = torch.nn.LayerNorm((dim), elementwise_affine=True)
+
+    def forward(self, x):
+        x = torch.swapaxes(x, -2, -1)
+        x = self.norm(x)
+        x = torch.swapaxes(x, -2, -1)
+        return x
+    
+class RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.g = nn.Parameter(torch.ones(1, 1, dim))
+
+    def forward(self, x):
+        return F.normalize(x, dim = -1) * self.g * (x.shape[-1] ** 0.5) #normalize divides by maximum norm element. Different from original in which we take the max norma nd not the sum of square elem. 
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = RMSNorm(dim)
+
+    def forward(self, x):
+        x = self.norm(x)
+        return self.fn(x)
+    
+    
+class Block(nn.Module):
+    def __init__(self, dim, dim_out, norm_type='none', act_type='tanh', *args, **kwargs):
+        super().__init__()
+        self.proj = StaticGraphLinear(dim, dim_out, *args, **kwargs)
+                                            # num_nodes=num_nodes,
+                                            # node_types=T)
+        if norm_type == 'none':
+            self.norm = nn.Identity() #nn.GroupNorm(groups, dim_out)
+        elif norm_type == 'layer':
+            self.norm = LayerNorm(kwargs['num_nodes'])
+        else: 
+            assert 0, f"Norm type {norm_type} not implemented!"
+        if act_type == 'tanh':
+            self.act = nn.Tanh()
+        else: 
+            assert 0, f"Activation type {act_type} not implemented!"
+
+    def forward(self, x, scale_shift = None):
+        x = self.proj(x)
+        x = self.norm(x)
+
+        if scale_shift is not None:
+            scale, shift = scale_shift
+            x = x * (scale + 1) + shift
+
+        x = self.act(x)
+        return x
+    
+    
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, dim_out, *, time_emb_dim = None, groups = 8, **kwargs):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Tanh(),
+            nn.Linear(time_emb_dim, dim_out * 2)
+        ) if time_emb_dim is not None else None
+
+        self.block1 = Block(dim, dim_out, groups = groups, **kwargs)
+        self.block2 = Block(dim_out, dim_out, groups = groups, **kwargs)
+        self.res_linear = StaticGraphLinear(dim, dim_out, bias=False, **kwargs)  if dim != dim_out else nn.Identity()
+
+    def forward(self, x, time_emb = None):
+
+        scale_shift = None
+        if self.mlp is not None and time_emb is not None:
+            time_emb = self.mlp(time_emb)
+            time_emb = rearrange(time_emb, 'b c -> b 1 c')
+            scale_shift = time_emb.chunk(2, dim = -1)
+
+        h = self.block1(x, scale_shift = scale_shift)
+
+        h = self.block2(h)
+
+        return h + self.res_linear(x)
+
+# We need default num_heads: int = 8,
+class Attention(nn.Module):
+    def __init__(self, dim, dim_out=None, heads = 4, dim_head = 32,qkv_bias: bool = False, attn_dropout: float = 0., proj_dropout: float = 0., qk_norm: bool = False, norm_layer: nn.Module = nn.Identity, **kwargs):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        dim_out = dim_out if dim_out is not None else dim
+
+        self.to_qkv = StaticGraphLinear(dim,hidden_dim * 3,bias=qkv_bias, **kwargs)
+        self.to_out = StaticGraphLinear(hidden_dim,dim_out,bias=False,**kwargs)
+        self.attn_dropout = nn.Dropout(attn_dropout)
+        self.out_dropout = nn.Dropout(proj_dropout)
+
+
+        self.q_norm = norm_layer(dim_head) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(dim_head) if qk_norm else nn.Identity()
+
+    def forward(self, x):
+        b, n, c = x.shape
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h c)  -> b h c n', h = self.heads), qkv)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        q = q * self.scale
+        sim = einsum('b h c n, b h c j -> b h n j', q, k)
+        attn = sim.softmax(dim = -1)
+        attn = self.attn_dropout(attn)
+
+        out = einsum('b h n j, b h d j -> b h n d', attn, v)
+
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.out_dropout(self.to_out(out))
+    
+    

SkeletonDiffusion/src/core/network/layers/graph_structural.py

@@ -0,0 +1,133 @@
+from typing import Tuple, Optional, List, Union
+
+import torch
+from torch.nn import *
+import math
+
+def gmm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
+    return torch.einsum('ndo,bnd->bno', w, x)
+
+
+class GraphLinear(Module):
+    def __init__(self, in_features: int, out_features: int):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+
+    def reset_parameters(self) -> None:
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        #stdv = 1. / math.sqrt(self.weight.size(1))
+        #self.weight.data.uniform_(-stdv, stdv)
+        #if self.learn_influence:
+        #    self.G.data.uniform_(-stdv, stdv)
+        if len(self.weight.shape) == 3:
+            self.weight.data[1:] = self.weight.data[0]
+        if self.bias is not None:
+            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in)
+            init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, input: torch.Tensor, g: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if g is None and self.learn_influence:
+            g = torch.nn.functional.normalize(self.G, p=1., dim=1)
+            #g = torch.softmax(self.G, dim=1)
+        elif g is None:
+            g = self.G
+        w = self.weight[self.node_type_index]
+        output = self.mm(input, w.transpose(-2, -1))
+        if self.bias is not None:
+            bias = self.bias[self.node_type_index]
+            output += bias
+        output = g.matmul(output)
+
+        return output
+
+
+class DynamicGraphLinear(GraphLinear):
+    def __init__(self, num_node_types: int = 1, *args):
+        super().__init__(*args)
+
+    def forward(self, input: torch.Tensor, g: torch.Tensor = None, t: torch.Tensor = None) -> torch.Tensor:
+        assert g is not None or t is not None, "Either Graph Influence Matrix or Node Type Vector is needed"
+        if g is None:
+            g = self.G[t][:, t]
+        return super().forward(input, g)
+
+
+
+class StaticGraphLinear(GraphLinear):
+    def __init__(self, *args, bias: bool = True, num_nodes: int = None, graph_influence: Union[torch.Tensor, Parameter] = None,
+                 learn_influence: bool = False, node_types: torch.Tensor = None, weights_per_type: bool = False, **kwargs):
+        """
+        :param in_features: Size of each input sample
+        :param out_features: Size of each output sample
+        :param num_nodes: Number of nodes.
+        :param graph_influence: Graph Influence Matrix
+        :param learn_influence: If set to ``False``, the layer will not learn an the Graph Influence Matrix.
+        :param node_types: List of Type for each node. All nodes of same type will share weights.
+                Default: All nodes have unique types.
+        :param weights_per_type: If set to ``False``, the layer will not learn weights for each node type.
+        :param bias: If set to ``False``, the layer will not learn an additive bias.
+        """
+        super().__init__(*args)
+
+        self.learn_influence = learn_influence
+
+        if graph_influence is not None:
+            assert num_nodes == graph_influence.shape[0] or num_nodes is None, 'Number of Nodes or Graph Influence Matrix has to be given.'
+            num_nodes = graph_influence.shape[0]
+            if type(graph_influence) is Parameter:
+                assert learn_influence, "Graph Influence Matrix is a Parameter, therefore it must be learnable."
+                self.G = graph_influence
+            elif learn_influence:
+                self.G = Parameter(graph_influence)
+            else:
+                self.register_buffer('G', graph_influence)
+        else:
+            assert num_nodes, 'Number of Nodes or Graph Influence Matrix has to be given.'
+            eye_influence = torch.eye(num_nodes, num_nodes)
+            if learn_influence:
+                self.G = Parameter(eye_influence)
+            else:
+                self.register_buffer('G', eye_influence)
+
+        if weights_per_type and node_types is None:
+            node_types = torch.tensor([i for i in range(num_nodes)])
+        if node_types is not None:
+            num_node_types = node_types.max() + 1
+            self.weight = Parameter(torch.Tensor(num_node_types, self.out_features, self.in_features))
+            self.mm = gmm
+            self.node_type_index = node_types
+        else:
+            self.weight = Parameter(torch.Tensor(self.out_features, self.in_features))
+            self.mm = torch.matmul
+            self.node_type_index = None
+
+        if bias:
+            if node_types is not None:
+                self.bias = Parameter(torch.Tensor(num_node_types, self.out_features))
+            else:
+                self.bias = Parameter(torch.Tensor(self.out_features))
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+
+# class BN(Module):
+#     def __init__(self, num_nodes, num_features):
+#         super().__init__()
+#         self.num_nodes = num_nodes
+#         self.num_features = num_features
+#         self.bn = BatchNorm1d(num_nodes * num_features)
+
+#     def forward(self, x: torch.Tensor) -> torch.Tensor:
+#         return self.bn(x.view(-1, self.num_nodes * self.num_features)).view(-1, self.num_nodes, self.num_features)
+
+# class LinearX(Module):
+#     def __init__(self):
+#         super().__init__()
+
+#     def forward(self, input: torch.Tensor) -> torch.Tensor:
+#         return input
+

SkeletonDiffusion/src/core/network/layers/recurrent.py

@@ -0,0 +1,402 @@
+from typing import Tuple, Optional, List, Union
+
+import torch
+from torch.nn import *
+import math
+
+from .graph_structural import gmm
+
+IFDEF_JITSCRIPT = False # generates nan for long sequences.
+
+GraphLSTMState = Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]
+
+class StaticGraphLSTMCell_(Module):
+    def __init__(self, input_size: int, hidden_size: int, num_nodes: int = None, dropout: float = 0.,
+                 recurrent_dropout: float = 0., graph_influence: Union[torch.Tensor, Parameter] = None,
+                 learn_influence: bool = False, additive_graph_influence: Union[torch.Tensor, Parameter] = None,
+                 learn_additive_graph_influence: bool = False, node_types: torch.Tensor = None,
+                 weights_per_type: bool = False, clockwork: bool = False, bias: bool = True):
+        """
+
+        :param input_size: The number of expected features in the input `x`
+        :param hidden_size: The number of features in the hidden state `h`
+        :param num_nodes:
+        :param dropout:
+        :param recurrent_dropout:
+        :param graph_influence:
+        :param learn_influence:
+        :param additive_graph_influence:
+        :param learn_additive_graph_influence:
+        :param node_types:
+        :param weights_per_type:
+        :param bias:
+        """
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+
+        self.learn_influence = learn_influence
+        self.learn_additive_graph_influence = learn_additive_graph_influence
+        if graph_influence is not None:
+            assert num_nodes == graph_influence.shape[0] or num_nodes is None, 'Number of Nodes or Graph Influence Matrix has to be given.'
+            num_nodes = graph_influence.shape[0]
+            if type(graph_influence) is Parameter:
+                assert learn_influence, "Graph Influence Matrix is a Parameter, therefore it must be learnable."
+                self.G = graph_influence
+            elif learn_influence:
+                self.G = Parameter(graph_influence)
+            else:
+                self.register_buffer('G', graph_influence)
+        else:
+            assert num_nodes, 'Number of Nodes or Graph Influence Matrix has to be given.'
+            eye_influence = torch.eye(num_nodes, num_nodes)
+            if learn_influence:
+                self.G = Parameter(eye_influence)
+            else:
+                self.register_buffer('G', eye_influence)
+
+        if additive_graph_influence is not None:
+            if type(additive_graph_influence) is Parameter:
+                self.G_add = additive_graph_influence
+            elif learn_additive_graph_influence:
+                self.G_add = Parameter(additive_graph_influence)
+            else:
+                self.register_buffer('G_add', additive_graph_influence)
+        else:
+            if learn_additive_graph_influence:
+                self.G_add = Parameter(torch.zeros_like(self.G))
+            else:
+                self.G_add = 0.
+
+        if weights_per_type and node_types is None:
+            node_types = torch.tensor([i for i in range(num_nodes)])
+        if node_types is not None:
+            num_node_types = node_types.max() + 1
+            self.weight_ih = Parameter(torch.Tensor(num_node_types, 4 * hidden_size, input_size))
+            self.weight_hh = Parameter(torch.Tensor(num_node_types, 4 * hidden_size, hidden_size))
+            self.mm = gmm
+            self.register_buffer('node_type_index', node_types)
+        else:
+            self.weight_ih = Parameter(torch.Tensor(4 * hidden_size, input_size))
+            self.weight_hh = Parameter(torch.Tensor(4 * hidden_size, hidden_size))
+            self.mm = torch.matmul
+            self.register_buffer('node_type_index', None)
+
+        if bias:
+            if node_types is not None:
+                self.bias_ih = Parameter(torch.Tensor(num_node_types, 4 * hidden_size))
+                self.bias_hh = Parameter(torch.Tensor(num_node_types, 4 * hidden_size))
+            else:
+                self.bias_ih = Parameter(torch.Tensor(4 * hidden_size))
+                self.bias_hh = Parameter(torch.Tensor(4 * hidden_size))
+        else:
+            self.bias_ih = None
+            self.bias_hh = None
+
+        self.clockwork = clockwork
+        if clockwork:
+            phase = torch.arange(0., hidden_size)
+            phase = phase - phase.min()
+            phase = (phase / phase.max()) * 8.
+            phase += 1.
+            phase = torch.floor(phase)
+            self.register_buffer('phase', phase)
+        else:
+            phase = torch.ones(hidden_size)
+            self.register_buffer('phase', phase)
+
+        self.dropout = Dropout(dropout)
+        self.r_dropout = Dropout(recurrent_dropout)
+
+        self.num_nodes = num_nodes
+
+        self.init_weights()
+
+    def init_weights(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            if weight is self.G:
+                continue
+            if weight is self.G_add:
+                continue
+            weight.data.uniform_(-stdv, stdv)
+            if weight is self.weight_hh or weight is self.weight_ih and len(self.weight_ih.shape) == 3:
+                weight.data[1:] = weight.data[0]
+
+    def forward(self, input: torch.Tensor, state: GraphLSTMState, t: int = 0) -> Tuple[torch.Tensor, GraphLSTMState]:
+        hx, cx, gx = state
+        if hx is None:
+            hx = torch.zeros(input.shape[0], self.num_nodes, self.hidden_size, dtype=input.dtype, device=input.device)
+        if cx is None:
+            cx = torch.zeros(input.shape[0], self.num_nodes, self.hidden_size, dtype=input.dtype, device=input.device)
+        if gx is None and self.learn_influence:
+            gx = torch.nn.functional.normalize(self.G, p=1., dim=1)
+            #gx = torch.softmax(self.G, dim=1)
+        elif gx is None:
+            gx = self.G
+
+        hx = self.r_dropout(hx)
+
+        weight_ih = self.weight_ih[self.node_type_index]
+        weight_hh = self.weight_hh[self.node_type_index]
+        if self.bias_hh is not None:
+            bias_hh = self.bias_hh[self.node_type_index]
+        else:
+            bias_hh = 0.
+
+        c_mask = (torch.remainder(torch.tensor(t + 1., device=input.device), self.phase) < 0.01).type_as(cx)
+
+        gates = (self.dropout(self.mm(input, weight_ih.transpose(-2, -1))) +
+                 self.mm(hx, weight_hh.transpose(-2, -1)) + bias_hh)
+        gates = torch.matmul(gx, gates)
+        ingate, forgetgate, cellgate, outgate = gates.chunk(4, 2)
+
+        ingate = torch.sigmoid(ingate)
+        forgetgate = torch.sigmoid(forgetgate)
+        cellgate = torch.tanh(cellgate)
+        outgate = torch.sigmoid(outgate)
+
+        cy = c_mask * ((forgetgate * cx) + (ingate * cellgate)) + (1 - c_mask) * cx
+        hy = outgate * torch.tanh(cy)
+
+        gx = gx + self.G_add
+        if self.learn_influence or self.learn_additive_graph_influence:
+            gx = torch.nn.functional.normalize(gx, p=1., dim=1)
+            #gx = torch.softmax(gx, dim=1)
+
+        return hy, (hy, cy, gx)
+
+
+class StaticGraphLSTM_(Module):
+    def __init__(self, input_size: int, hidden_size: int, num_layers: int = 1, layer_dropout: float = 0.0, **kwargs):
+        super().__init__()
+        self.layers = ModuleList([StaticGraphLSTMCell_(input_size, hidden_size, **kwargs)]
+                                 + [StaticGraphLSTMCell_(hidden_size, hidden_size, **kwargs) for _ in range(num_layers - 1)])
+        self.dropout = Dropout(layer_dropout)
+
+    def forward(self, input: torch.Tensor, states: Optional[List[GraphLSTMState]] = None, t_i: int = 0) -> Tuple[torch.Tensor, List[GraphLSTMState]]:
+        if states is None:
+            n: Optional[torch.Tensor] = None
+            states = [(n, n, n)] * len(self.layers)
+
+        output_states: List[GraphLSTMState] = []
+        output = input
+        i = 0
+        for rnn_layer in self.layers:
+            state = states[i]
+            inputs = output.unbind(1)
+            outputs: List[torch.Tensor] = []
+            for t, input in enumerate(inputs):
+                out, state = rnn_layer(input, state, t_i+t)
+                outputs += [out]
+            output = torch.stack(outputs, dim=1)
+            output = self.dropout(output)
+            output_states += [state]
+            i += 1
+        return output, output_states
+
+
+def StaticGraphLSTM(*args, **kwargs):
+    if IFDEF_JITSCRIPT:
+        return torch.jit.script(StaticGraphLSTM_(*args, **kwargs))
+    else: 
+        return StaticGraphLSTM_(*args, **kwargs)
+
+GraphGRUState = Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]
+
+
+class StaticGraphGRUCell_(Module):
+    def __init__(self, input_size: int, hidden_size: int, num_nodes: int = None, dropout: float = 0.,
+                 recurrent_dropout: float = 0., graph_influence: Union[torch.Tensor, Parameter] = None,
+                 learn_influence: bool = False, additive_graph_influence: Union[torch.Tensor, Parameter] = None,
+                 learn_additive_graph_influence: bool = False, node_types: torch.Tensor = None,
+                 weights_per_type: bool = False, clockwork: bool = False, bias: bool = True):
+        """
+
+        :param input_size: The number of expected features in the input `x`
+        :param hidden_size: The number of features in the hidden state `h`
+        :param num_nodes:
+        :param dropout:
+        :param recurrent_dropout:
+        :param graph_influence:
+        :param learn_influence:
+        :param additive_graph_influence:
+        :param learn_additive_graph_influence:
+        :param node_types:
+        :param weights_per_type:
+        :param bias:
+        """
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+
+        self.learn_influence = learn_influence
+        self.learn_additive_graph_influence = learn_additive_graph_influence
+        if graph_influence is not None:
+            assert num_nodes == graph_influence.shape[0] or num_nodes is None, 'Number of Nodes or Graph Influence Matrix has to be given.'
+            num_nodes = graph_influence.shape[0]
+            if type(graph_influence) is Parameter:
+                assert learn_influence, "Graph Influence Matrix is a Parameter, therefore it must be learnable."
+                self.G = graph_influence
+            elif learn_influence:
+                self.G = Parameter(graph_influence)
+            else:
+                self.register_buffer('G', graph_influence)
+        else:
+            assert num_nodes, 'Number of Nodes or Graph Influence Matrix has to be given.'
+            eye_influence = torch.eye(num_nodes, num_nodes)
+            if learn_influence:
+                self.G = Parameter(eye_influence)
+            else:
+                self.register_buffer('G', eye_influence)
+
+        if additive_graph_influence is not None:
+            if type(additive_graph_influence) is Parameter:
+                self.G_add = additive_graph_influence
+            elif learn_additive_graph_influence:
+                self.G_add = Parameter(additive_graph_influence)
+            else:
+                self.register_buffer('G_add', additive_graph_influence)
+        else:
+            if learn_additive_graph_influence:
+                self.G_add = Parameter(torch.zeros_like(self.G))
+            else:
+                self.G_add = 0.
+
+        if weights_per_type and node_types is None:
+            node_types = torch.tensor([i for i in range(num_nodes)])
+        if node_types is not None:
+            num_node_types = node_types.max() + 1
+            self.weight_ih = Parameter(torch.Tensor(num_node_types, 3 * hidden_size, input_size))
+            self.weight_hh = Parameter(torch.Tensor(num_node_types, 3 * hidden_size, hidden_size))
+            self.mm = gmm
+            self.register_buffer('node_type_index', node_types)
+        else:
+            self.weight_ih = Parameter(torch.Tensor(3 * hidden_size, input_size))
+            self.weight_hh = Parameter(torch.Tensor(3 * hidden_size, hidden_size))
+            self.mm = torch.matmul
+            self.register_buffer('node_type_index', None)
+
+        if bias:
+            if node_types is not None:
+                self.bias_ih = Parameter(torch.Tensor(num_node_types, 3 * hidden_size))
+                self.bias_hh = Parameter(torch.Tensor(num_node_types, 3 * hidden_size))
+            else:
+                self.bias_ih = Parameter(torch.Tensor(3 * hidden_size))
+                self.bias_hh = Parameter(torch.Tensor(3 * hidden_size))
+        else:
+            self.bias_ih = None
+            self.bias_hh = None
+
+        self.clockwork = clockwork
+        if clockwork:
+            phase = torch.arange(0., hidden_size)
+            phase = phase - phase.min()
+            phase = (phase / phase.max()) * 8.
+            phase += 1.
+            phase = torch.floor(phase)
+            self.register_buffer('phase', phase)
+        else:
+            phase = torch.ones(hidden_size)
+            self.register_buffer('phase', phase)
+
+        self.dropout = Dropout(dropout)
+        self.r_dropout = Dropout(recurrent_dropout)
+
+        self.num_nodes = num_nodes
+
+        self.init_weights()
+
+    def init_weights(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            if weight is self.G:
+                continue
+            if weight is self.G_add:
+                continue
+            weight.data.uniform_(-stdv, stdv)
+            #if weight is self.weight_hh or weight is self.weight_ih and len(self.weight_ih.shape) == 3:
+            #    weight.data[1:] = weight.data[0]
+
+    def forward(self, input: torch.Tensor, state: GraphGRUState, t: int = 0) -> Tuple[torch.Tensor, GraphGRUState]:
+        hx, gx = state
+        if hx is None:
+            hx = torch.zeros(input.shape[0], self.num_nodes, self.hidden_size, dtype=input.dtype, device=input.device)
+        if gx is None and self.learn_influence:
+            gx = torch.nn.functional.normalize(self.G, p=1., dim=1)
+            #gx = torch.softmax(self.G, dim=1)
+        elif gx is None:
+            gx = self.G
+
+        hx = self.r_dropout(hx)
+
+        weight_ih = self.weight_ih[self.node_type_index]
+        weight_hh = self.weight_hh[self.node_type_index]
+        if self.bias_hh is not None:
+            bias_hh = self.bias_hh[self.node_type_index]
+        else:
+            bias_hh = 0.
+        if self.bias_ih is not None:
+            bias_ih = self.bias_ih[self.node_type_index]
+        else:
+            bias_ih = 0.
+
+        c_mask = (torch.remainder(torch.tensor(t + 1., device=input.device), self.phase) < 0.01).type_as(hx)
+
+        x_results = self.dropout(self.mm(input, weight_ih.transpose(-2, -1))) + bias_ih
+        h_results = self.mm(hx, weight_hh.transpose(-2, -1)) + bias_hh
+        x_results = torch.matmul(gx, x_results)
+        h_results = torch.matmul(gx, h_results)
+
+        i_r, i_z, i_n = x_results.chunk(3, 2)
+        h_r, h_z, h_n = h_results.chunk(3, 2)
+
+        r = torch.sigmoid(i_r + h_r)
+        z = torch.sigmoid(i_z + h_z)
+        n = torch.tanh(i_n + r * h_n)
+
+        hy = n - torch.mul(n, z) + torch.mul(z, hx)
+        hy = c_mask * hy + (1 - c_mask) * hx
+
+        gx = gx + self.G_add
+        if self.learn_influence or self.learn_additive_graph_influence:
+            gx = torch.nn.functional.normalize(gx, p=1., dim=1)
+            #gx = torch.softmax(gx, dim=1)
+
+        return hy, (hy, gx)
+
+
+class StaticGraphGRU_(Module):
+    def __init__(self, input_size: int, hidden_size: int, num_layers: int = 1, layer_dropout: float = 0.0, **kwargs):
+        super().__init__()
+        self.layers = ModuleList([StaticGraphGRUCell_(input_size, hidden_size, **kwargs)]
+                                 + [StaticGraphGRUCell_(hidden_size, hidden_size, **kwargs) for _ in range(num_layers - 1)])
+        self.dropout = Dropout(layer_dropout)
+
+    def forward(self, input: torch.Tensor, states: Optional[List[GraphGRUState]] = None, t_i: int = 0) -> Tuple[torch.Tensor, List[GraphGRUState]]:
+        if states is None:
+            n: Optional[torch.Tensor] = None
+            states = [(n, n)] * len(self.layers)
+
+        output_states: List[GraphGRUState] = []
+        output = input
+        i = 0
+        for rnn_layer in self.layers:
+            state = states[i]
+            inputs = output.unbind(1)
+            outputs: List[torch.Tensor] = []
+            for t, input in enumerate(inputs):
+                out, state = rnn_layer(input, state, t_i+t)
+                outputs += [out]
+            output = torch.stack(outputs, dim=1)
+            output = self.dropout(output)
+            output_states += [state]
+            i += 1
+        return output, output_states
+
+
+def StaticGraphGRU(*args, **kwargs):
+    if IFDEF_JITSCRIPT:
+        return torch.jit.script(StaticGraphGRU_(*args, **kwargs))
+    else: 
+        return StaticGraphGRU_(*args, **kwargs)

SkeletonDiffusion/src/core/network/nn/__init__.py

@@ -0,0 +1,2 @@
+from .generator import Denoiser
+from .autoencoder import AutoEncoder

SkeletonDiffusion/src/core/network/nn/autoencoder.py

@@ -0,0 +1,105 @@
+from typing import Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from ..layers import StaticGraphLinear
+from .decoder import Decoder
+from .encoder import Encoder
+
+        
+class AutoEncoder(nn.Module):
+    def __init__(self,
+                num_nodes: int,
+                encoder_hidden_size: int,
+                decoder_hidden_size: int,
+                latent_size: int,
+                node_types: torch.Tensor = None,
+                input_size: int = 3,
+                z_activation: str = 'tanh',
+                enc_num_layers: int = 1,
+                loss_pose_type: str = 'l1',
+                **kwargs):
+        super().__init__()
+        self.param_groups = [{}]
+        self.latent_size = latent_size
+        self.loss_pose_type = loss_pose_type 
+
+        self.encoder = Encoder(num_nodes=num_nodes,
+                               input_size=input_size,
+                               hidden_size=encoder_hidden_size,
+                               output_size=latent_size,
+                               node_types=node_types,
+                               enc_num_layers = enc_num_layers,
+                               recurrent_arch = kwargs['recurrent_arch_enc'],)
+
+        
+        assert kwargs['output_size'] == input_size
+        self.decoder = Decoder( num_nodes=num_nodes,
+                        input_size=latent_size ,
+                        feature_size=input_size,
+                        hidden_size=decoder_hidden_size,
+                        node_types=node_types,
+                        param_groups=self.param_groups,
+                        **kwargs
+                        )
+        assert z_activation in ['tanh', 'identity'], f"z_activation must be either 'tanh' or 'identity', but got {z_activation}"
+        self.z_activation = nn.Tanh() if z_activation == "tanh" else nn.Identity()
+
+    
+    def forward(self, x):
+        h, _ = self.encoder(x)
+        return h
+    
+    def get_past_embedding(self, past, state=None):
+        with torch.no_grad():
+            h_hat_embedding = self(past)
+        z_past = self.z_activation(h_hat_embedding)
+        return z_past
+    
+    def get_embedding(self, future, state=None):
+        z = self.forward(future)
+        return z
+    
+    def get_train_embeddings(self, y, past, state=None):
+        z_past = self.get_past_embedding(past, state=state)
+        z = self.get_embedding(y, state=state)
+        return z_past, z
+    
+    def decode(self, x: torch.Tensor, h: torch.Tensor, z: torch.Tensor, ph=1, state=None):
+        x_tiled = x[:, -2:] 
+        out, _ = self.decoder(x=x_tiled,
+                                h=h,
+                                z=z,
+                                ph=ph,
+                                state=state)  # [B * Z, T, N, D]
+        return out
+
+    def autoencode(self, y, past, ph=1, state=None):
+        z_past, z = self.get_train_embeddings(y, past, state=state)
+        out = self.decode(past, z, z_past, ph)
+        return out, z_past, z
+
+    def loss(self, y_pred, y, type=None, reduction="mean", **kwargs):
+        type = self.loss_pose_type if type is None else type
+        if type=="mse":
+            out = torch.nn.MSELoss(reduction="none")(y_pred,y)
+        elif type in ["l1", "L1"]:
+            out = torch.nn.L1Loss(reduction="none")(y_pred,y)
+        else: 
+            assert 0, "Not implemnted"
+        loss = (out.sum(-1) #spatial size 
+                    .mean(-1) #keypoints
+                    .mean(-1) # timesteps
+                    )
+        if reduction == "mean":
+            return loss.mean()
+        elif reduction == "none":
+            return loss
+        else: 
+            assert 0, "Not implemnted"
+        return loss
+
+    
+    
+