Upload 242 files

.gitattributes

@@ -33,3 +33,4 @@ 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
+difpoint/src/utils/dependencies/insightface/data/images/t1.jpg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,124 @@
+import os, sys
+import gradio as gr
+from difpoint.inference import Inferencer
+from TTS.api import TTS
+import torch
+import time
+from flask import send_from_directory
+from huggingface_hub import snapshot_download
+import spaces
+import tensorrt
+import multiprocessing as mp
+import pickle
+mp.set_start_method('spawn', force=True)  
+
+
+repo_id = "ChaolongYang/KDTalker"
+local_dir = "./downloaded_repo"  
+snapshot_download(repo_id=repo_id, local_dir=local_dir)
+print("\nFiles downloaded:")
+for root, dirs, files in os.walk(local_dir):
+    for file in files:
+        file_path = os.path.join(root, file)
+        print(file_path)
+
+result_dir = "results"
+def set_upload():
+    return "upload"
+def set_microphone():
+    return "microphone"
+def set_tts():
+    return "tts"
+def create_kd_talker():
+    return Inferencer() 
+
+@spaces.GPU
+def predict(prompt, upload_reference_audio, microphone_reference_audio, reference_audio_type):
+    global result_dir
+    output_file_path = os.path.join('./downloaded_repo/', 'output.wav')
+    if reference_audio_type == 'upload':
+        audio_file_pth = upload_reference_audio
+    elif reference_audio_type == 'microphone':
+        audio_file_pth =  microphone_reference_audio
+    tts = TTS('tts_models/multilingual/multi-dataset/your_tts')
+    tts.tts_to_file(
+        text=prompt,
+        file_path=output_file_path,
+        speaker_wav=audio_file_pth,
+        language="en",
+    )
+    return gr.Audio(value=output_file_path, type='filepath')
+
+@spaces.GPU
+def generate(upload_driven_audio, tts_driven_audio, driven_audio_type, source_image, smoothed_pitch, smoothed_yaw, smoothed_roll, smoothed_t):
+    return Inferencer().generate_with_audio_img(upload_driven_audio, tts_driven_audio, driven_audio_type, source_image,
+                                    smoothed_pitch, smoothed_yaw, smoothed_roll, smoothed_t)
+
+
+def main():
+    if torch.cuda.is_available():
+        device = "cuda" 
+    else:
+        device = "cpu"
+
+    with gr.Blocks(analytics_enabled=False) as interface:
+        gr.Markdown(
+        """
+            <div align='center'>
+                <h2> Unlock Pose Diversity: Accurate and Efficient Implicit Keypoint-based Spatiotemporal Diffusion for Audio-driven Talking Portrait </h2>
+                <div style="display: flex; justify-content: center; align-items: center; gap: 20px;">
+                    <img src='https://newstatic.dukekunshan.edu.cn/mainsite/2021/08/07161629/large_dku-Logo-e1649298929570.png' alt='Logo' width='150'/>
+                    <img src='https://www.xjtlu.edu.cn/wp-content/uploads/2023/12/7c52fd62e9cf26cb493faa7f91c2782.png' width='250'/>
+                </div>
+            </div>
+        """
+        )
+        driven_audio_type = gr.Textbox(value="upload", visible=False) 
+        reference_audio_type = gr.Textbox(value="upload", visible=False)
+
+        with gr.Row():
+            with gr.Column(variant="panel"):
+                with gr.Tabs(elem_id="kdtalker_source_image"):
+                    with gr.TabItem("Upload image"):
+                        source_image = gr.Image(label="Source image", sources="upload", type="filepath", scale=256)
+                
+                with gr.Tabs(elem_id="kdtalker_driven_audio"):
+                    with gr.TabItem("Upload"):
+                        upload_driven_audio = gr.Audio(label="Upload audio", sources="upload", type="filepath")
+                        upload_driven_audio.change(set_upload, outputs=driven_audio_type)
+                    with gr.TabItem("TTS"):
+                        upload_reference_audio = gr.Audio(label="Upload Reference Audio", sources="upload", type="filepath")
+                        upload_reference_audio.change(set_upload, outputs=reference_audio_type)
+                        microphone_reference_audio = gr.Audio(label="Recorded Reference Audio", sources="microphone", type="filepath")
+                        microphone_reference_audio.change(set_microphone, outputs=reference_audio_type)
+                        input_text = gr.Textbox(
+                            label="Generating audio from text",
+                            lines=5,
+                            placeholder="please enter some text here, we generate the audio from text using @Coqui.ai TTS."
+                        )
+                        tts_button = gr.Button("Generate audio", elem_id="kdtalker_audio_generate", variant="primary")
+                        tts_driven_audio = gr.Audio(label="Synthesised Audio", type="filepath")
+                        tts_button.click(fn=predict, inputs=[input_text, upload_reference_audio, microphone_reference_audio, reference_audio_type], outputs=[tts_driven_audio])
+                        tts_button.click(set_tts, outputs=driven_audio_type)
+            with gr.Column(variant="panel"):
+                gen_video = gr.Video(label="Generated video", format="mp4", width=256)
+                with gr.Tabs(elem_id="talker_checkbox"):
+                    with gr.TabItem("KDTalker"):
+                        smoothed_pitch = gr.Slider(minimum=0, maximum=1, step=0.1, label="Pitch", value=0.8)
+                        smoothed_yaw = gr.Slider(minimum=0, maximum=1, step=0.1, label="Yaw", value=0.8)
+                        smoothed_roll = gr.Slider(minimum=0, maximum=1, step=0.1, label="Roll", value=0.8)
+                        smoothed_t = gr.Slider(minimum=0, maximum=1, step=0.1, label="T", value=0.8)
+                        kd_submit = gr.Button("Generate", elem_id="kdtalker_generate", variant="primary")
+                        kd_submit.click(
+                                fn=generate,
+                                inputs=[
+                                    upload_driven_audio, tts_driven_audio, driven_audio_type, source_image,
+                                    smoothed_pitch, smoothed_yaw, smoothed_roll, smoothed_t
+                                ],
+                                outputs=[gen_video]
+                            )
+    return interface
+
+
+demo = main()
+demo.queue().launch()

difpoint/configs/onnx_infer.yaml

@@ -0,0 +1,105 @@
+models:
+  warping_spade:
+    name: "WarpingSpadeModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/warping_spade.onnx"
+  motion_extractor:
+    name: "MotionExtractorModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/motion_extractor.onnx"
+  landmark:
+    name: "LandmarkModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/landmark.onnx"
+  face_analysis:
+    name: "FaceAnalysisModel"
+    predict_type: "ort"
+    model_path:
+      - "./downloaded_repo/pretrained_weights/liveportrait_onnx/retinaface_det_static.onnx"
+      - "./downloaded_repo/pretrained_weights/liveportrait_onnx/face_2dpose_106_static.onnx"
+  app_feat_extractor:
+    name: "AppearanceFeatureExtractorModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/appearance_feature_extractor.onnx"
+  stitching:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching.onnx"
+  stitching_eye_retarget:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching_eye.onnx"
+  stitching_lip_retarget:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching_lip.onnx"
+
+animal_models:
+  warping_spade:
+    name: "WarpingSpadeModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/warping_spade.onnx"
+  motion_extractor:
+    name: "MotionExtractorModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/motion_extractor.onnx"
+  app_feat_extractor:
+    name: "AppearanceFeatureExtractorModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/appearance_feature_extractor.onnx"
+  stitching:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching.onnx"
+  stitching_eye_retarget:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching_eye.onnx"
+  stitching_lip_retarget:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching_lip.onnx"
+  landmark:
+    name: "LandmarkModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/landmark.onnx"
+  face_analysis:
+    name: "FaceAnalysisModel"
+    predict_type: "ort"
+    model_path:
+      - "./downloaded_repo/pretrained_weights/liveportrait_onnx/retinaface_det_static.onnx"
+      - "./downloaded_repo/pretrained_weights/liveportrait_onnx/face_2dpose_106_static.onnx"
+
+crop_params:
+  src_dsize: 512
+  src_scale: 2.3
+  src_vx_ratio: 0.0
+  src_vy_ratio: -0.125
+  dri_scale: 2.2
+  dri_vx_ratio: 0.0
+  dri_vy_ratio: -0.1
+
+
+infer_params:
+  flag_crop_driving_video: False
+  flag_normalize_lip: True
+  flag_source_video_eye_retargeting: False
+  flag_video_editing_head_rotation: False
+  flag_eye_retargeting: False
+  flag_lip_retargeting: False
+  flag_stitching: True
+  flag_relative_motion: True
+  flag_pasteback: True
+  flag_do_crop: True
+  flag_do_rot: True
+
+  # NOT EXPOERTED PARAMS
+  lip_normalize_threshold: 0.03 # threshold for flag_normalize_lip
+  source_video_eye_retargeting_threshold: 0.18 # threshold for eyes retargeting if the input is a source video
+  driving_smooth_observation_variance: 1e-7 # smooth strength scalar for the animated video when the input is a source video, the larger the number, the smoother the animated video; too much smoothness would result in loss of motion accuracy
+  anchor_frame: 0 # TO IMPLEMENT
+  mask_crop_path: "./assets/mask_template.png"
+  driving_multiplier: 1.0
+
+  source_max_dim: 1280 # the max dim of height and width of source image
+  source_division: 2 # make sure the height and width of source image can be divided by this number

difpoint/configs/onnx_mp_infer.yaml

@@ -0,0 +1,99 @@
+models:
+  warping_spade:
+    name: "WarpingSpadeModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/warping_spade.onnx"
+  motion_extractor:
+    name: "MotionExtractorModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/motion_extractor.onnx"
+  landmark:
+    name: "LandmarkModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/landmark.onnx"
+  face_analysis:
+    name: "MediaPipeFaceModel"
+    predict_type: "mp"
+  app_feat_extractor:
+    name: "AppearanceFeatureExtractorModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/appearance_feature_extractor.onnx"
+  stitching:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching.onnx"
+  stitching_eye_retarget:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching_eye.onnx"
+  stitching_lip_retarget:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching_lip.onnx"
+
+animal_models:
+  warping_spade:
+    name: "WarpingSpadeModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/warping_spade.onnx"
+  motion_extractor:
+    name: "MotionExtractorModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/motion_extractor.onnx"
+  app_feat_extractor:
+    name: "AppearanceFeatureExtractorModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/appearance_feature_extractor.onnx"
+  stitching:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching.onnx"
+  stitching_eye_retarget:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching_eye.onnx"
+  stitching_lip_retarget:
+    name: "StitchingModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching_lip.onnx"
+  landmark:
+    name: "LandmarkModel"
+    predict_type: "ort"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/landmark.onnx"
+  face_analysis:
+    name: "MediaPipeFaceModel"
+    predict_type: "mp"
+
+crop_params:
+  src_dsize: 512
+  src_scale: 2.3
+  src_vx_ratio: 0.0
+  src_vy_ratio: -0.125
+  dri_scale: 2.2
+  dri_vx_ratio: 0.0
+  dri_vy_ratio: -0.1
+
+
+infer_params:
+  flag_crop_driving_video: False
+  flag_normalize_lip: True
+  flag_source_video_eye_retargeting: False
+  flag_video_editing_head_rotation: False
+  flag_eye_retargeting: False
+  flag_lip_retargeting: False
+  flag_stitching: True
+  flag_relative_motion: True
+  flag_pasteback: True
+  flag_do_crop: True
+  flag_do_rot: True
+
+  # NOT EXPOERTED PARAMS
+  lip_normalize_threshold: 0.03 # threshold for flag_normalize_lip
+  source_video_eye_retargeting_threshold: 0.18 # threshold for eyes retargeting if the input is a source video
+  driving_smooth_observation_variance: 1e-7 # smooth strength scalar for the animated video when the input is a source video, the larger the number, the smoother the animated video; too much smoothness would result in loss of motion accuracy
+  anchor_frame: 0 # TO IMPLEMENT
+  mask_crop_path: "./assets/mask_template.png"
+  driving_multiplier: 1.0
+
+  source_max_dim: 1280 # the max dim of height and width of source image
+  source_division: 2 # make sure the height and width of source image can be divided by this number

difpoint/configs/trt_infer.yaml

@@ -0,0 +1,105 @@
+models:
+  warping_spade:
+    name: "WarpingSpadeModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/warping_spade-fix.trt"
+  motion_extractor:
+    name: "MotionExtractorModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/motion_extractor.trt"
+  landmark:
+    name: "LandmarkModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/landmark.trt"
+  face_analysis:
+    name: "FaceAnalysisModel"
+    predict_type: "trt"
+    model_path:
+      - "./downloaded_repo/pretrained_weights/liveportrait_onnx/retinaface_det_static.trt"
+      - "./downloaded_repo/pretrained_weights/liveportrait_onnx/face_2dpose_106_static.trt"
+  app_feat_extractor:
+    name: "AppearanceFeatureExtractorModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/appearance_feature_extractor.trt"
+  stitching:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching.trt"
+  stitching_eye_retarget:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching_eye.trt"
+  stitching_lip_retarget:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching_lip.trt"
+
+animal_models:
+  warping_spade:
+    name: "WarpingSpadeModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/warping_spade-fix.trt"
+  motion_extractor:
+    name: "MotionExtractorModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/motion_extractor.trt"
+  app_feat_extractor:
+    name: "AppearanceFeatureExtractorModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/appearance_feature_extractor.trt"
+  stitching:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching.trt"
+  stitching_eye_retarget:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching_eye.trt"
+  stitching_lip_retarget:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching_lip.trt"
+  landmark:
+    name: "LandmarkModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/landmark.trt"
+  face_analysis:
+    name: "FaceAnalysisModel"
+    predict_type: "trt"
+    model_path:
+      - "./downloaded_repo/pretrained_weights/liveportrait_onnx/retinaface_det_static.trt"
+      - "./downloaded_repo/pretrained_weights/liveportrait_onnx/face_2dpose_106_static.trt"
+
+crop_params:
+  src_dsize: 512
+  src_scale: 2.3
+  src_vx_ratio: 0.0
+  src_vy_ratio: -0.125
+  dri_scale: 2.2
+  dri_vx_ratio: 0.0
+  dri_vy_ratio: -0.1
+
+
+infer_params:
+  flag_crop_driving_video: False
+  flag_normalize_lip: False
+  flag_source_video_eye_retargeting: False
+  flag_video_editing_head_rotation: False
+  flag_eye_retargeting: False
+  flag_lip_retargeting: False
+  flag_stitching: True
+  flag_relative_motion: True
+  flag_pasteback: True
+  flag_do_crop: True
+  flag_do_rot: True
+
+  # NOT EXPOERTED PARAMS
+  lip_normalize_threshold: 0.03 # threshold for flag_normalize_lip
+  source_video_eye_retargeting_threshold: 0.18 # threshold for eyes retargeting if the input is a source video
+  driving_smooth_observation_variance: 1e-7 # smooth strength scalar for the animated video when the input is a source video, the larger the number, the smoother the animated video; too much smoothness would result in loss of motion accuracy
+  anchor_frame: 0 # TO IMPLEMENT
+  mask_crop_path: "./assets/mask_template.png"
+  driving_multiplier: 1.0
+
+  source_max_dim: 1280 # the max dim of height and width of source image
+  source_division: 2 # make sure the height and width of source image can be divided by this number

difpoint/configs/trt_mp_infer.yaml

@@ -0,0 +1,99 @@
+models:
+  warping_spade:
+    name: "WarpingSpadeModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/warping_spade-fix.trt"
+  motion_extractor:
+    name: "MotionExtractorModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/motion_extractor.trt"
+  landmark:
+    name: "LandmarkModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/landmark.trt"
+  face_analysis:
+    name: "MediaPipeFaceModel"
+    predict_type: "mp"
+  app_feat_extractor:
+    name: "AppearanceFeatureExtractorModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/appearance_feature_extractor.trt"
+  stitching:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching.trt"
+  stitching_eye_retarget:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching_eye.trt"
+  stitching_lip_retarget:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/stitching_lip.trt"
+
+animal_models:
+  warping_spade:
+    name: "WarpingSpadeModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/warping_spade-fix.trt"
+  motion_extractor:
+    name: "MotionExtractorModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/motion_extractor.trt"
+  app_feat_extractor:
+    name: "AppearanceFeatureExtractorModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/appearance_feature_extractor.trt"
+  stitching:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching.trt"
+  stitching_eye_retarget:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching_eye.trt"
+  stitching_lip_retarget:
+    name: "StitchingModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_animal_onnx/stitching_lip.trt"
+  landmark:
+    name: "LandmarkModel"
+    predict_type: "trt"
+    model_path: "./downloaded_repo/pretrained_weights/liveportrait_onnx/landmark.trt"
+  face_analysis:
+    name: "MediaPipeFaceModel"
+    predict_type: "mp"
+
+crop_params:
+  src_dsize: 512
+  src_scale: 2.0
+  src_vx_ratio: 0.0
+  src_vy_ratio: -0.125
+  dri_scale: 2.2
+  dri_vx_ratio: 0.0
+  dri_vy_ratio: -0.1
+
+
+infer_params:
+  flag_crop_driving_video: False
+  flag_normalize_lip: False
+  flag_source_video_eye_retargeting: False
+  flag_video_editing_head_rotation: False
+  flag_eye_retargeting: False
+  flag_lip_retargeting: False
+  flag_stitching: False
+  flag_relative_motion: False
+  flag_pasteback: False
+  flag_do_crop: False
+  flag_do_rot: False
+
+  # NOT EXPOERTED PARAMS
+  lip_normalize_threshold: 0.03 # threshold for flag_normalize_lip
+  source_video_eye_retargeting_threshold: 0.18 # threshold for eyes retargeting if the input is a source video
+  driving_smooth_observation_variance: 1e-7 # smooth strength scalar for the animated video when the input is a source video, the larger the number, the smoother the animated video; too much smoothness would result in loss of motion accuracy
+  anchor_frame: 0 # TO IMPLEMENT
+  mask_crop_path: "./assets/mask_template.png"
+  driving_multiplier: 1.0
+
+  source_max_dim: 1280 # the max dim of height and width of source image
+  source_division: 2 # make sure the height and width of source image can be divided by this number

difpoint/croper.py

@@ -0,0 +1,269 @@
+import os
+import cv2
+import time
+import glob
+import argparse
+import scipy
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from itertools import cycle
+
+
+"""
+brief: face alignment with FFHQ method (https://github.com/NVlabs/ffhq-dataset)
+author: lzhbrian (https://lzhbrian.me)
+date: 2020.1.5
+note: code is heavily borrowed from 
+    https://github.com/NVlabs/ffhq-dataset
+    http://dlib.net/face_landmark_detection.py.html
+requirements:
+    apt install cmake
+    conda install Pillow numpy scipy
+    pip install dlib
+    # download face landmark model from: 
+    # http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
+"""
+
+import numpy as np
+from PIL import Image
+import dlib
+
+
+class Croper:
+    
+    def __init__(self, path_of_lm):
+        # download model from: http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
+        self.predictor = dlib.shape_predictor(path_of_lm)
+
+    def get_landmark(self, img_np):
+        """get landmark with dlib
+        :return: np.array shape=(68, 2)
+        """
+        detector = dlib.get_frontal_face_detector()
+        dets = detector(img_np, 1)
+        #     print("Number of faces detected: {}".format(len(dets)))
+        #     for k, d in enumerate(dets):
+        if len(dets) == 0:
+            return None
+        d = dets[0]
+        # Get the landmarks/parts for the face in box d.
+        shape = self.predictor(img_np, d)
+        #         print("Part 0: {}, Part 1: {} ...".format(shape.part(0), shape.part(1)))
+        t = list(shape.parts())
+        a = []
+        for tt in t:
+            a.append([tt.x, tt.y])
+        lm = np.array(a)
+        # lm is a shape=(68,2) np.array
+        return lm
+
+    def align_face(self, img, lm, output_size=1024):
+        """
+        :param filepath: str
+        :return: PIL Image
+        """
+        lm_chin = lm[0: 17]  # left-right
+        lm_eyebrow_left = lm[17: 22]  # left-right
+        lm_eyebrow_right = lm[22: 27]  # left-right
+        lm_nose = lm[27: 31]  # top-down
+        lm_nostrils = lm[31: 36]  # top-down
+        lm_eye_left = lm[36: 42]  # left-clockwise
+        lm_eye_right = lm[42: 48]  # left-clockwise
+        lm_mouth_outer = lm[48: 60]  # left-clockwise
+        lm_mouth_inner = lm[60: 68]  # left-clockwise
+
+        # Calculate auxiliary vectors.
+        eye_left = np.mean(lm_eye_left, axis=0)
+        eye_right = np.mean(lm_eye_right, axis=0)
+        eye_avg = (eye_left + eye_right) * 0.5
+        eye_to_eye = eye_right - eye_left
+        mouth_left = lm_mouth_outer[0]
+        mouth_right = lm_mouth_outer[6]
+        mouth_avg = (mouth_left + mouth_right) * 0.5
+        eye_to_mouth = mouth_avg - eye_avg
+
+        # Choose oriented crop rectangle.
+        x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]  # Addition of binocular difference and double mouth difference
+        x /= np.hypot(*x)   # hypot函数计算直角三角形的斜边长，用斜边长对三角形两条直边做归一化
+        x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)    # 双眼差和眼嘴差，选较大的作为基准尺度
+        y = np.flipud(x) * [-1, 1]
+        c = eye_avg + eye_to_mouth * 0.1
+        quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])   # 定义四边形，以面部基准位置为中心上下左右平移得到四个顶点
+        qsize = np.hypot(*x) * 2    # 定义四边形的大小（边长），为基准尺度的2倍
+
+        # Shrink.
+        # 如果计算出的四边形太大了，就按比例缩小它
+        shrink = int(np.floor(qsize / output_size * 0.5))
+        if shrink > 1:
+            rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink)))
+            img = img.resize(rsize, Image.ANTIALIAS)
+            quad /= shrink
+            qsize /= shrink
+        else:
+            rsize = (int(np.rint(float(img.size[0]))), int(np.rint(float(img.size[1]))))
+
+        # Crop.
+        border = max(int(np.rint(qsize * 0.1)), 3)
+        crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+                int(np.ceil(max(quad[:, 1]))))
+        crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]),
+                min(crop[3] + border, img.size[1]))
+        if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]:
+            # img = img.crop(crop)
+            quad -= crop[0:2]
+
+        # Pad.
+        pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+               int(np.ceil(max(quad[:, 1]))))
+        pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0),
+               max(pad[3] - img.size[1] + border, 0))
+        # if enable_padding and max(pad) > border - 4:
+        #     pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
+        #     img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
+        #     h, w, _ = img.shape
+        #     y, x, _ = np.ogrid[:h, :w, :1]
+        #     mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w - 1 - x) / pad[2]),
+        #                       1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h - 1 - y) / pad[3]))
+        #     blur = qsize * 0.02
+        #     img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
+        #     img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)
+        #     img = Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB')
+        #     quad += pad[:2]
+
+        # Transform.
+        quad = (quad + 0.5).flatten()
+        lx = max(min(quad[0], quad[2]), 0)
+        ly = max(min(quad[1], quad[7]), 0)
+        rx = min(max(quad[4], quad[6]), img.size[0])
+        ry = min(max(quad[3], quad[5]), img.size[0])
+        # img = img.transform((transform_size, transform_size), Image.QUAD, (quad + 0.5).flatten(),
+        #                     Image.BILINEAR)
+        # if output_size < transform_size:
+        #     img = img.resize((output_size, output_size), Image.ANTIALIAS)
+
+        # Save aligned image.
+        return rsize, crop, [lx, ly, rx, ry]
+
+    # def crop(self, img_np_list):
+    #     for _i in range(len(img_np_list)):
+    #         img_np = img_np_list[_i]
+    #         lm = self.get_landmark(img_np)
+    #         if lm is None:
+    #             return None
+    #         crop, quad = self.align_face(img=Image.fromarray(img_np), lm=lm, output_size=512)
+    #         clx, cly, crx, cry = crop
+    #         lx, ly, rx, ry = quad
+    #         lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry)
+        
+    #         _inp = img_np_list[_i]
+    #         _inp = _inp[cly:cry, clx:crx]
+    #         _inp = _inp[ly:ry, lx:rx]
+    #         img_np_list[_i] = _inp
+    #     return img_np_list
+    
+    def crop(self, img_np_list, still=False, xsize=512):    # first frame for all video
+        img_np = img_np_list[0]
+        lm = self.get_landmark(img_np)
+        if lm is None:
+            raise 'can not detect the landmark from source image'
+        rsize, crop, quad = self.align_face(img=Image.fromarray(img_np), lm=lm, output_size=xsize)
+        clx, cly, crx, cry = crop
+        lx, ly, rx, ry = quad
+        lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry)
+        for _i in range(len(img_np_list)):
+            _inp = img_np_list[_i]
+            _inp = cv2.resize(_inp, (rsize[0], rsize[1]))
+            _inp = _inp[cly:cry, clx:crx]
+            # cv2.imwrite('test1.jpg', _inp)
+            if not still:
+                _inp = _inp[ly:ry, lx:rx]
+            # cv2.imwrite('test2.jpg', _inp)
+            img_np_list[_i] = _inp
+        return img_np_list, crop, quad
+
+
+def read_video(filename, uplimit=100):
+    frames = []
+    cap = cv2.VideoCapture(filename)
+    cnt = 0
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if ret:
+            frame = cv2.resize(frame, (512, 512))
+            frames.append(frame)
+        else:
+            break
+        cnt += 1
+        if cnt >= uplimit:
+            break
+    cap.release()
+    assert len(frames) > 0, f'{filename}: video with no frames!'
+    return frames
+
+
+def create_video(video_name, frames, fps=25, video_format='.mp4', resize_ratio=1):
+    # video_name = os.path.dirname(image_folder) + video_format
+    # img_list = glob.glob1(image_folder, 'frame*')
+    # img_list.sort()
+    # frame = cv2.imread(os.path.join(image_folder, img_list[0]))
+    # frame = cv2.resize(frame, (0, 0), fx=resize_ratio, fy=resize_ratio)
+    # height, width, layers = frames[0].shape
+    height, width, layers = 512, 512, 3
+    if video_format == '.mp4':
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    elif video_format == '.avi':
+        fourcc = cv2.VideoWriter_fourcc(*'XVID')
+    video = cv2.VideoWriter(video_name, fourcc, fps, (width, height))
+    for _frame in frames:
+        _frame = cv2.resize(_frame, (height, width), interpolation=cv2.INTER_LINEAR)
+        video.write(_frame)
+
+def create_images(video_name, frames):
+    height, width, layers = 512, 512, 3
+    images_dir = video_name.split('.')[0]
+    os.makedirs(images_dir, exist_ok=True)
+    for i, _frame in enumerate(frames):
+        _frame = cv2.resize(_frame, (height, width), interpolation=cv2.INTER_LINEAR)
+        _frame_path = os.path.join(images_dir, str(i)+'.jpg')
+        cv2.imwrite(_frame_path, _frame)
+
+def run(data):
+    filename, opt, device = data
+    os.environ['CUDA_VISIBLE_DEVICES'] = device
+    croper = Croper()
+
+    frames = read_video(filename, uplimit=opt.uplimit)
+    name = filename.split('/')[-1]  # .split('.')[0]
+    name = os.path.join(opt.output_dir, name)
+
+    frames = croper.crop(frames)
+    if frames is None:
+        print(f'{name}: detect no face. should removed')
+        return
+    # create_video(name, frames)
+    create_images(name, frames)
+
+
+def get_data_path(video_dir):
+    eg_video_files = ['/apdcephfs/share_1290939/quincheng/datasets/HDTF/backup_fps25/WDA_KatieHill_000.mp4']
+    # filenames = list()
+    # VIDEO_EXTENSIONS_LOWERCASE = {'mp4'}
+    # VIDEO_EXTENSIONS = VIDEO_EXTENSIONS_LOWERCASE.union({f.upper() for f in VIDEO_EXTENSIONS_LOWERCASE})
+    # extensions = VIDEO_EXTENSIONS
+    # for ext in extensions:
+    #     filenames = sorted(glob.glob(f'{opt.input_dir}/**/*.{ext}'))
+    # print('Total number of videos:', len(filenames))
+    return eg_video_files
+
+
+def get_wra_data_path(video_dir):
+    if opt.option == 'video':
+        videos_path = sorted(glob.glob(f'{video_dir}/*.mp4'))
+    elif opt.option == 'image':
+        videos_path = sorted(glob.glob(f'{video_dir}/*/'))
+    else:
+        raise NotImplementedError
+    print('Example videos: ', videos_path[:2])
+    return videos_path
+

difpoint/dataset_process/audio.py

@@ -0,0 +1,156 @@
+import librosa
+import librosa.filters
+import numpy as np
+# import tensorflow as tf
+from scipy import signal
+from scipy.io import wavfile
+from difpoint.src.utils.hparams import hparams as hp
+
+
+def load_wav(path, sr):
+    return librosa.core.load(path, sr=sr)[0]
+
+
+def save_wav(wav, path, sr):
+    wav *= 32767 / max(0.01, np.max(np.abs(wav)))
+    # proposed by @dsmiller
+    wavfile.write(path, sr, wav.astype(np.int16))
+
+
+def save_wavenet_wav(wav, path, sr):
+    librosa.output.write_wav(path, wav, sr=sr)
+
+
+def preemphasis(wav, k, preemphasize=True):
+    if preemphasize:
+        return signal.lfilter([1, -k], [1], wav)
+    return wav
+
+
+def inv_preemphasis(wav, k, inv_preemphasize=True):
+    if inv_preemphasize:
+        return signal.lfilter([1], [1, -k], wav)
+    return wav
+
+
+def get_hop_size():
+    hop_size = hp.hop_size
+    if hop_size is None:
+        assert hp.frame_shift_ms is not None
+        hop_size = int(hp.frame_shift_ms / 1000 * hp.sample_rate)
+    return hop_size
+
+
+def linearspectrogram(wav):
+    D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
+    S = _amp_to_db(np.abs(D)) - hp.ref_level_db
+
+    if hp.signal_normalization:
+        return _normalize(S)
+    return S
+
+
+def melspectrogram(wav):
+    D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
+    S = _amp_to_db(_linear_to_mel(np.abs(D))) - hp.ref_level_db
+
+    if hp.signal_normalization:
+        return _normalize(S)
+    return S
+
+
+def _lws_processor():
+    import lws
+    return lws.lws(hp.n_fft, get_hop_size(), fftsize=hp.win_size, mode="speech")
+
+
+def _stft(y):
+    if hp.use_lws:
+        return _lws_processor(hp).stft(y).T
+    else:
+        return librosa.stft(y=y, n_fft=hp.n_fft, hop_length=get_hop_size(), win_length=hp.win_size)
+
+
+##########################################################
+# Those are only correct when using lws!!! (This was messing with Wavenet quality for a long time!)
+def num_frames(length, fsize, fshift):
+    """Compute number of time frames of spectrogram
+    """
+    pad = (fsize - fshift)
+    if length % fshift == 0:
+        M = (length + pad * 2 - fsize) // fshift + 1
+    else:
+        M = (length + pad * 2 - fsize) // fshift + 2
+    return M
+
+
+def pad_lr(x, fsize, fshift):
+    """Compute left and right padding
+    """
+    M = num_frames(len(x), fsize, fshift)
+    pad = (fsize - fshift)
+    T = len(x) + 2 * pad
+    r = (M - 1) * fshift + fsize - T
+    return pad, pad + r
+
+
+##########################################################
+# Librosa correct padding
+def librosa_pad_lr(x, fsize, fshift):
+    return 0, (x.shape[0] // fshift + 1) * fshift - x.shape[0]
+
+
+# Conversions
+_mel_basis = None
+
+
+def _linear_to_mel(spectogram):
+    global _mel_basis
+    if _mel_basis is None:
+        _mel_basis = _build_mel_basis()
+    return np.dot(_mel_basis, spectogram)
+
+
+def _build_mel_basis():
+    assert hp.fmax <= hp.sample_rate // 2
+    return librosa.filters.mel(sr=hp.sample_rate, n_fft=hp.n_fft, n_mels=hp.num_mels,
+                               fmin=hp.fmin, fmax=hp.fmax)
+
+
+def _amp_to_db(x):
+    min_level = np.exp(hp.min_level_db / 20 * np.log(10))
+    return 20 * np.log10(np.maximum(min_level, x))
+
+
+def _db_to_amp(x):
+    return np.power(10.0, (x) * 0.05)
+
+
+def _normalize(S):
+    if hp.allow_clipping_in_normalization:
+        if hp.symmetric_mels:
+            return np.clip((2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value,
+                           -hp.max_abs_value, hp.max_abs_value)
+        else:
+            return np.clip(hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db)), 0, hp.max_abs_value)
+
+    assert S.max() <= 0 and S.min() - hp.min_level_db >= 0
+    if hp.symmetric_mels:
+        return (2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value
+    else:
+        return hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db))
+
+
+def _denormalize(D):
+    if hp.allow_clipping_in_normalization:
+        if hp.symmetric_mels:
+            return (((np.clip(D, -hp.max_abs_value,
+                              hp.max_abs_value) + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value))
+                    + hp.min_level_db)
+        else:
+            return ((np.clip(D, 0, hp.max_abs_value) * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db)
+
+    if hp.symmetric_mels:
+        return (((D + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value)) + hp.min_level_db)
+    else:
+        return ((D * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db)

difpoint/dataset_process/wav2lip.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Conv2d(nn.Module):
+    def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, use_act=True, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.conv_block = nn.Sequential(
+            nn.Conv2d(cin, cout, kernel_size, stride, padding),
+            nn.BatchNorm2d(cout)
+        )
+        self.act = nn.ReLU()
+        self.residual = residual
+        self.use_act = use_act
+
+    def forward(self, x):
+        out = self.conv_block(x)
+        if self.residual:
+            out += x
+
+        if self.use_act:
+            return self.act(out)
+        else:
+            return out
+
+class AudioEncoder(nn.Module):
+    def __init__(self, wav2lip_checkpoint, device):
+        super(AudioEncoder, self).__init__()
+
+        self.audio_encoder = nn.Sequential(
+            Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
+            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
+            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
+
+            Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),
+            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
+            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
+
+            Conv2d(64, 128, kernel_size=3, stride=3, padding=1),
+            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
+            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
+
+            Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),
+            Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
+
+            Conv2d(256, 512, kernel_size=3, stride=1, padding=0),
+            Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)
+
+        #### load the pre-trained audio_encoder
+        wav2lip_state_dict = torch.load(wav2lip_checkpoint, map_location=torch.device(device))['state_dict']
+        state_dict = self.audio_encoder.state_dict()
+
+        for k,v in wav2lip_state_dict.items():
+            if 'audio_encoder' in k:
+                state_dict[k.replace('module.audio_encoder.', '')] = v
+        self.audio_encoder.load_state_dict(state_dict)
+
+
+    def forward(self, audio_sequences):
+        # audio_sequences = (B, T, 1, 80, 16)
+        B = audio_sequences.size(0)
+
+        audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)
+
+        audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 1
+        dim = audio_embedding.shape[1]
+        audio_embedding = audio_embedding.reshape((B, -1, dim, 1, 1))
+
+        return audio_embedding.squeeze(-1).squeeze(-1) #B seq_len+1 512
+
+wav2lip_checkpoint='ckpts/wav2lip.pth'
+wav2lip_model = AudioEncoder(wav2lip_checkpoint, 'cuda')
+wav2lip_model.cuda()
+wav2lip_model.eval()

difpoint/datasets/norm_info_d6.5_c8.5_vox1_train.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9422e503e75df9d1bd455d8e0f9f5e2826b12956cdedbb5566097c0151bddafb
+size 5580

difpoint/inference.py

@@ -0,0 +1,474 @@
+# -*- coding: UTF-8 -*-
+'''
+@File    ：inference.py
+@Author  ：Chaolong Yang
+@Date    ：2024/5/29 19:26 
+'''
+import glob
+
+import os
+os.environ['HYDRA_FULL_ERROR']='1'
+
+
+import os
+import time
+import shutil
+import uuid
+import os
+import cv2
+import tyro
+
+import numpy as np
+from tqdm import tqdm
+import cv2
+from rich.progress import track
+
+from difpoint.croper import Croper
+from PIL import Image
+import time
+
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+import imageio
+from pydub import AudioSegment
+from pykalman import KalmanFilter
+import scipy
+import matplotlib.pyplot as plt
+import matplotlib
+matplotlib.use('Agg')
+
+from difpoint.dataset_process import audio
+import os
+import argparse
+import pdb
+import ffmpeg
+import cv2
+import time
+import numpy as np
+import os
+import datetime
+import platform
+from omegaconf import OmegaConf
+#from difpoint.src.pipelines.faster_live_portrait_pipeline import FasterLivePortraitPipeline
+from difpoint.src.live_portrait_pipeline import LivePortraitPipeline
+from difpoint.src.config.argument_config import ArgumentConfig
+from difpoint.src.config.inference_config import InferenceConfig
+from difpoint.src.config.crop_config import CropConfig
+from difpoint.src.live_portrait_pipeline import LivePortraitPipeline
+from difpoint.src.utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio
+from difpoint.src.utils.camera import get_rotation_matrix
+from difpoint.src.utils.video import images2video, concat_frames, get_fps, add_audio_to_video, has_audio_stream
+
+
+FFMPEG = "ffmpeg"
+
+def parse_audio_length(audio_length, sr, fps):
+    bit_per_frames = sr / fps
+    num_frames = int(audio_length / bit_per_frames)
+    audio_length = int(num_frames * bit_per_frames)
+    return audio_length, num_frames
+
+def crop_pad_audio(wav, audio_length):
+    if len(wav) > audio_length:
+        wav = wav[:audio_length]
+    elif len(wav) < audio_length:
+        wav = np.pad(wav, [0, audio_length - len(wav)], mode='constant', constant_values=0)
+    return wav
+
+class Conv2d(nn.Module):
+    def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, use_act=True, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.conv_block = nn.Sequential(
+            nn.Conv2d(cin, cout, kernel_size, stride, padding),
+            nn.BatchNorm2d(cout)
+        )
+        self.act = nn.ReLU()
+        self.residual = residual
+        self.use_act = use_act
+
+    def forward(self, x):
+        out = self.conv_block(x)
+        if self.residual:
+            out += x
+
+        if self.use_act:
+            return self.act(out)
+        else:
+            return out
+
+class AudioEncoder(nn.Module):
+    def __init__(self, wav2lip_checkpoint, device):
+        super(AudioEncoder, self).__init__()
+
+        self.audio_encoder = nn.Sequential(
+            Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
+            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
+            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
+
+            Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),
+            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
+            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
+
+            Conv2d(64, 128, kernel_size=3, stride=3, padding=1),
+            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
+            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
+
+            Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),
+            Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
+
+            Conv2d(256, 512, kernel_size=3, stride=1, padding=0),
+            Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)
+
+        #### load the pre-trained audio_encoder
+        wav2lip_state_dict = torch.load(wav2lip_checkpoint, map_location=torch.device(device))['state_dict']
+        state_dict = self.audio_encoder.state_dict()
+
+        for k,v in wav2lip_state_dict.items():
+            if 'audio_encoder' in k:
+                state_dict[k.replace('module.audio_encoder.', '')] = v
+        self.audio_encoder.load_state_dict(state_dict)
+
+    def forward(self, audio_sequences):
+        # audio_sequences = (B, T, 1, 80, 16)
+        B = audio_sequences.size(0)
+
+        audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)
+
+        audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 1
+        dim = audio_embedding.shape[1]
+        audio_embedding = audio_embedding.reshape((B, -1, dim, 1, 1))
+
+        return audio_embedding.squeeze(-1).squeeze(-1) #B seq_len+1 512
+
+def partial_fields(target_class, kwargs):
+    return target_class(**{k: v for k, v in kwargs.items() if hasattr(target_class, k)})
+
+def dct2device(dct: dict, device):
+    for key in dct:
+        dct[key] = torch.tensor(dct[key]).to(device)
+    return dct
+
+def save_video_with_watermark(video, audio, save_path, watermark=False):
+    temp_file = str(uuid.uuid4())+'.mp4'
+    cmd = r'ffmpeg -y -i "%s" -i "%s" -vcodec copy "%s"' % (video, audio, temp_file)
+    os.system(cmd)
+    shutil.move(temp_file, save_path)
+    
+
+    
+class Inferencer(object):
+    def __init__(self):
+        
+        st=time.time()
+        print('#'*25+'Start initialization'+'#'*25)
+        self.device = 'cuda'
+        from difpoint.model import get_model
+        self.point_diffusion = get_model()
+        ckpt = torch.load('./downloaded_repo/ckpts/KDTalker.pth', weights_only=False)
+
+        self.point_diffusion.load_state_dict(ckpt['model'])
+        print('model', self.point_diffusion.children())
+        self.point_diffusion.eval()
+        self.point_diffusion.to(self.device)
+
+        lm_croper_checkpoint = './downloaded_repo/ckpts/shape_predictor_68_face_landmarks.dat'
+        self.croper = Croper(lm_croper_checkpoint)
+
+        self.norm_info = dict(np.load(r'difpoint/datasets/norm_info_d6.5_c8.5_vox1_train.npz'))
+
+        wav2lip_checkpoint = './downloaded_repo/ckpts/wav2lip.pth'
+        self.wav2lip_model = AudioEncoder(wav2lip_checkpoint, 'cuda')
+        self.wav2lip_model.cuda()
+        self.wav2lip_model.eval()
+
+        args = tyro.cli(ArgumentConfig)
+
+        self.inf_cfg = partial_fields(InferenceConfig, args.__dict__)  # use attribute of args to initial InferenceConfig
+        self.crop_cfg = partial_fields(CropConfig, args.__dict__)  # use attribute of args to initial CropConfig
+
+        self.live_portrait_pipeline = LivePortraitPipeline(inference_cfg=self.inf_cfg, crop_cfg=self.crop_cfg)
+        print('#'*25+f'End initialization, cost time {time.time()-st}'+'#'*25)
+
+    def _norm(self, data_dict):
+        for k in data_dict.keys():
+            if k in ['yaw', 'pitch', 'roll', 't', 'exp', 'scale', 'kp', 'c_lip', 'c_eye']:
+                v=data_dict[k]
+                data_dict[k] = (v - self.norm_info[k+'_mean'])/self.norm_info[k+'_std']
+        return data_dict
+
+    def _denorm(self, data_dict):
+        for k in data_dict.keys():
+            if k in ['yaw', 'pitch', 'roll', 't', 'exp', 'scale', 'kp', 'c_lip', 'c_eye']:
+                v=data_dict[k]
+                data_dict[k] = v * self.norm_info[k+'_std'] + self.norm_info[k+'_mean']
+        return data_dict
+
+
+    def output_to_dict(self, data):
+        output = {}
+
+        output['scale'] = data[:, 0]
+        output['yaw'] = data[:, 1, None]
+        output['pitch'] = data[:, 2, None]
+        output['roll'] = data[:, 3, None]
+        output['t'] = data[:, 4:7]
+        output['exp'] = data[:, 7:]
+
+        return output
+
+    def extract_mel_from_audio(self, audio_file_path):
+        syncnet_mel_step_size = 16
+        fps = 25
+        wav = audio.load_wav(audio_file_path, 16000)
+        wav_length, num_frames = parse_audio_length(len(wav), 16000, 25)
+        wav = crop_pad_audio(wav, wav_length)
+        orig_mel = audio.melspectrogram(wav).T
+        spec = orig_mel.copy()
+        indiv_mels = []
+
+        for i in tqdm(range(num_frames), 'mel:'):
+            start_frame_num = i - 2
+            start_idx = int(80. * (start_frame_num / float(fps)))
+            end_idx = start_idx + syncnet_mel_step_size
+            seq = list(range(start_idx, end_idx))
+            seq = [min(max(item, 0), orig_mel.shape[0] - 1) for item in seq]
+            m = spec[seq, :]
+            indiv_mels.append(m.T)
+        indiv_mels = np.asarray(indiv_mels)  # T 80 16
+        return indiv_mels
+
+    def extract_wav2lip_from_audio(self, audio_file_path):
+        asd_mel = self.extract_mel_from_audio(audio_file_path)
+        asd_mel = torch.FloatTensor(asd_mel).cuda().unsqueeze(0).unsqueeze(2)
+        with torch.no_grad():
+            hidden = self.wav2lip_model(asd_mel)
+        return hidden[0].cpu().detach().numpy()
+
+    def headpose_pred_to_degree(self, pred):
+        device = pred.device
+        idx_tensor = [idx for idx in range(66)]
+        idx_tensor = torch.FloatTensor(idx_tensor).to(device)
+        pred = F.softmax(pred)
+        degree = torch.sum(pred * idx_tensor, 1) * 3 - 99
+        return degree
+
+    def calc_combined_eye_ratio(self, c_d_eyes_i, c_s_eyes):
+        c_s_eyes_tensor = torch.from_numpy(c_s_eyes).float().to(self.device)
+        c_d_eyes_i_tensor = c_d_eyes_i[0].reshape(1, 1).to(self.device)
+        # [c_s,eyes, c_d,eyes,i]
+        combined_eye_ratio_tensor = torch.cat([c_s_eyes_tensor, c_d_eyes_i_tensor], dim=1)
+        return combined_eye_ratio_tensor
+
+    def calc_combined_lip_ratio(self, c_d_lip_i, c_s_lip):
+        c_s_lip_tensor = torch.from_numpy(c_s_lip).float().to(self.device)
+        c_d_lip_i_tensor = c_d_lip_i[0].to(self.device).reshape(1, 1) # 1x1
+        # [c_s,lip, c_d,lip,i]
+        combined_lip_ratio_tensor = torch.cat([c_s_lip_tensor, c_d_lip_i_tensor], dim=1) # 1x2
+        return combined_lip_ratio_tensor
+
+    # 2024.06.26
+    @torch.no_grad()
+    def generate_with_audio_img(self, upload_audio_path, tts_audio_path, audio_type, image_path, smoothed_pitch, smoothed_yaw, smoothed_roll, smoothed_t,  save_path='./downloaded_repo/'):
+        print(audio_type)
+        if audio_type == 'upload':
+            audio_path = upload_audio_path
+        elif audio_type == 'tts':
+            audio_path = tts_audio_path
+        save_path = os.path.join(save_path, "output.mp4")
+        image = [np.array(Image.open(image_path).convert('RGB'))]
+        if image[0].shape[0] != 256 or image[0].shape[1] != 256:
+            cropped_image, crop, quad = self.croper.crop(image, still=False, xsize=512)
+            input_image = cv2.resize(cropped_image[0], (256, 256))
+        else:
+            input_image = image[0]
+
+        I_s = torch.FloatTensor(input_image.transpose((2, 0, 1))).unsqueeze(0).cuda() / 255
+
+        x_s_info = self.live_portrait_pipeline.live_portrait_wrapper.get_kp_info(I_s)
+        x_c_s = x_s_info['kp'].reshape(1, 21, -1)
+        R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
+        f_s = self.live_portrait_pipeline.live_portrait_wrapper.extract_feature_3d(I_s)
+        x_s = self.live_portrait_pipeline.live_portrait_wrapper.transform_keypoint(x_s_info)
+
+        flag_lip_zero = self.inf_cfg.flag_lip_zero  # not overwrite 
+
+
+        ######## process driving info ########
+        kp_info = {}
+        for k in x_s_info.keys():
+            kp_info[k] = x_s_info[k].cpu().numpy()
+        # kp_info['c_lip'] = c_s_lip
+        # kp_info['c_eye'] = c_s_eye
+
+        kp_info = self._norm(kp_info)
+
+        ori_kp = torch.cat([torch.zeros([1, 7]), torch.Tensor(kp_info['kp'])], -1).cuda()
+
+        input_x = np.concatenate([kp_info[k] for k in ['scale', 'yaw', 'pitch', 'roll', 't']], 1)
+        input_x = np.concatenate((input_x, kp_info['exp'].reshape(1, 63)), axis=1)
+        input_x = np.expand_dims(input_x, -1)
+        input_x = np.expand_dims(input_x, 0)
+        input_x = np.concatenate([input_x, input_x, input_x], -1)
+
+        aud_feat = self.extract_wav2lip_from_audio(audio_path)
+
+        outputs = [input_x]
+
+        st = time.time()
+        print('#' * 25 + 'Start Inference' + '#' * 25)
+        sample_frame = 64    # 32 aud_feat.shape[0]
+        
+        for i in range(0, aud_feat.shape[0] - 1, sample_frame):
+            input_mel = torch.Tensor(aud_feat[i: i + sample_frame]).unsqueeze(0).cuda()
+            kp0 = torch.Tensor(outputs[-1])[:, -1].cuda()
+            pred_kp = self.point_diffusion.forward_sample(70, ref_kps=kp0, ori_kps=ori_kp, aud_feat=input_mel,
+                                                          scheduler='ddim', num_inference_steps=50)
+            outputs.append(pred_kp.cpu().numpy())
+
+
+        outputs = np.mean(np.concatenate(outputs, 1)[0], -1)[1:, ]
+        output_dict = self.output_to_dict(outputs)
+        output_dict = self._denorm(output_dict)
+
+        num_frame = output_dict['yaw'].shape[0]
+        x_d_info = {}
+        for key in output_dict:
+            x_d_info[key] = torch.tensor(output_dict[key]).cuda()
+
+        # smooth
+        def smooth(sequence, n_dim_state=1):
+            kf = KalmanFilter(initial_state_mean=sequence[0],
+                              transition_covariance=0.05 * np.eye(n_dim_state),  # 较小的过程噪声
+                              observation_covariance=0.001 * np.eye(n_dim_state))  # 可以增大观测噪声，减少敏感性
+            state_means, _ = kf.smooth(sequence)
+            return state_means
+
+        # scale_data = x_d_info['scale'].cpu().numpy()
+        yaw_data = x_d_info['yaw'].cpu().numpy()
+        pitch_data = x_d_info['pitch'].cpu().numpy()
+        roll_data = x_d_info['roll'].cpu().numpy()
+        t_data = x_d_info['t'].cpu().numpy()
+        exp_data = x_d_info['exp'].cpu().numpy()
+
+        smoothed_pitch = smooth(pitch_data, n_dim_state=1) * smoothed_pitch
+        smoothed_yaw = smooth(yaw_data, n_dim_state=1) * smoothed_yaw
+        smoothed_roll = smooth(roll_data, n_dim_state=1) * smoothed_roll
+        # smoothed_scale = smooth(scale_data, n_dim_state=1)
+        smoothed_t = smooth(t_data, n_dim_state=3) * smoothed_t
+        smoothed_exp = smooth(exp_data, n_dim_state=63)
+
+        # x_d_info['scale'] = torch.Tensor(smoothed_scale).cuda()
+        x_d_info['pitch'] = torch.Tensor(smoothed_pitch).cuda()
+        x_d_info['yaw'] = torch.Tensor(smoothed_yaw).cuda()
+        x_d_info['roll'] = torch.Tensor(smoothed_roll).cuda()
+        x_d_info['t'] = torch.Tensor(smoothed_t).cuda()
+        x_d_info['exp'] = torch.Tensor(smoothed_exp).cuda()
+
+
+
+        template_dct = {'motion': [], 'c_d_eyes_lst': [], 'c_d_lip_lst': []}
+        for i in track(range(num_frame), description='Making motion templates...', total=num_frame):
+            # collect s_d, R_d, δ_d and t_d for inference
+            x_d_i_info = x_d_info
+            R_d_i = get_rotation_matrix(x_d_i_info['pitch'][i], x_d_i_info['yaw'][i], x_d_i_info['roll'][i])
+
+            item_dct = {
+                'scale': x_d_i_info['scale'][i].cpu().numpy().astype(np.float32),
+                'R_d': R_d_i.cpu().numpy().astype(np.float32),
+                'exp': x_d_i_info['exp'][i].reshape(1, 21, -1).cpu().numpy().astype(np.float32),
+                't': x_d_i_info['t'][i].cpu().numpy().astype(np.float32),
+            }
+
+            template_dct['motion'].append(item_dct)
+            # template_dct['c_d_eyes_lst'].append(x_d_i_info['c_eye'][i])
+            # template_dct['c_d_lip_lst'].append(x_d_i_info['c_lip'][i])
+
+        I_p_lst = []
+        R_d_0, x_d_0_info = None, None
+
+        for i in track(range(num_frame), description='Animating...', total=num_frame):
+            x_d_i_info = template_dct['motion'][i]
+
+            for key in x_d_i_info:
+                x_d_i_info[key] = torch.tensor(x_d_i_info[key]).cuda()
+
+            R_d_i = x_d_i_info['R_d']
+
+            if i == 0:
+                R_d_0 = R_d_i
+                x_d_0_info = x_d_i_info
+            
+
+            if self.inf_cfg.flag_relative_motion:
+                R_new = (R_d_i @ R_d_0.permute(0, 2, 1)) @ R_s
+                delta_new = x_s_info['exp'].reshape(1, 21, -1) + (x_d_i_info['exp'] - x_d_0_info['exp'])
+                scale_new = x_s_info['scale'] * (x_d_i_info['scale'] / x_d_0_info['scale'])
+                t_new = x_s_info['t'] + (x_d_i_info['t'] - x_d_0_info['t'])
+            else:
+                R_new = R_d_i
+                delta_new = x_d_i_info['exp']
+                scale_new = x_s_info['scale']
+                t_new = x_d_i_info['t']
+            t_new[..., 2] = 0  # zero tz
+            
+            x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new
+    
+
+            # Algorithm 1:
+            if not self.inf_cfg.flag_stitching and not self.inf_cfg.flag_eye_retargeting and not self.inf_cfg.flag_lip_retargeting:
+                # without stitching or retargeting
+                if flag_lip_zero:
+                    x_d_i_new += lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
+                else:
+                    pass
+            elif self.inf_cfg.flag_stitching and not self.inf_cfg.flag_eye_retargeting and not self.inf_cfg.flag_lip_retargeting:
+                # with stitching and without retargeting
+                if flag_lip_zero:
+                    x_d_i_new = self.live_portrait_pipeline.live_portrait_wrapper.stitching(x_s, x_d_i_new) + lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
+                else:
+                    x_d_i_new = self.live_portrait_pipeline.live_portrait_wrapper.stitching(x_s, x_d_i_new)
+            else:
+                eyes_delta, lip_delta = None, None
+
+                if self.inf_cfg.flag_relative_motion:  # use x_s
+                    x_d_i_new = x_s + \
+                                (eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
+                                (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
+                else:  # use x_d,i
+                    x_d_i_new = x_d_i_new + \
+                                (eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
+                                (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
+
+                if self.inf_cfg.flag_stitching:
+                    x_d_i_new = self.live_portrait_pipeline.live_portrait_wrapper.stitching(x_s, x_d_i_new)
+
+
+            out = self.live_portrait_pipeline.live_portrait_wrapper.warp_decode(f_s, x_s, x_d_i_new)
+            I_p_i = self.live_portrait_pipeline.live_portrait_wrapper.parse_output(out['out'])[0]
+            I_p_lst.append(I_p_i)
+
+        video_name = os.path.basename(save_path)
+        video_save_dir = os.path.dirname(save_path)
+        path = os.path.join(video_save_dir, video_name)
+                  
+        imageio.mimsave(path, I_p_lst, fps=float(25))
+
+        audio_name = audio_path.split('/')[-1]
+        new_audio_path = os.path.join(video_save_dir, audio_name)
+        start_time = 0
+        # cog will not keep the .mp3 filename
+        sound = AudioSegment.from_file(audio_path)
+        end_time = start_time + num_frame * 1 / 25 * 1000
+        word1 = sound.set_frame_rate(16000)
+        word = word1[start_time:end_time]
+        word.export(new_audio_path, format="wav")
+
+        save_video_with_watermark(path, new_audio_path, save_path, watermark=False)
+        print(f'The generated video is named {video_save_dir}/{video_name}')
+
+        print('#' * 25 + f'End Inference, cost time {time.time() - st}' + '#' * 25)
+        return save_path
+    
+    

difpoint/model/__init__.py

@@ -0,0 +1,6 @@
+from .model import ConditionalPointCloudDiffusionModel
+
+def get_model():
+    model = ConditionalPointCloudDiffusionModel()
+    return model
+

difpoint/model/model.py

@@ -0,0 +1,409 @@
+import inspect
+from typing import Optional
+from einops import rearrange
+import torch
+import torch.nn.functional as F
+from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
+from diffusers.schedulers.scheduling_ddim import DDIMScheduler
+from diffusers.schedulers.scheduling_pndm import PNDMScheduler
+
+from torch import Tensor
+from tqdm import tqdm
+from diffusers import ModelMixin
+from .model_utils import get_custom_betas
+from .point_model import PointModel
+
+import copy
+class ConditionalPointCloudDiffusionModel(ModelMixin):
+    def __init__(
+        self,
+        beta_start: float = 1e-5,
+        beta_end: float = 8e-3,
+        beta_schedule: str = 'linear',
+        point_cloud_model: str = 'simple',
+        point_cloud_model_embed_dim: int = 64,
+    ):
+        super().__init__()
+        self.in_channels = 70  # 3 for 3D point positions
+        self.out_channels = 70
+
+        # Checks
+        # Create diffusion model schedulers which define the sampling timesteps
+        scheduler_kwargs = {}
+        if beta_schedule == 'custom':
+            scheduler_kwargs.update(dict(trained_betas=get_custom_betas(beta_start=beta_start, beta_end=beta_end)))
+        else:
+            scheduler_kwargs.update(dict(beta_start=beta_start, beta_end=beta_end, beta_schedule=beta_schedule))
+        self.schedulers_map = {
+            'ddpm': DDPMScheduler(**scheduler_kwargs, clip_sample=False),
+            'ddim': DDIMScheduler(**scheduler_kwargs, clip_sample=False), 
+            'pndm': PNDMScheduler(**scheduler_kwargs), 
+        }
+        self.scheduler = self.schedulers_map['ddim']  # this can be changed for inference
+
+        # Create point cloud model for processing point cloud at each diffusion step
+        self.point_model = PointModel(
+            model_type=point_cloud_model,
+            embed_dim=point_cloud_model_embed_dim,
+            in_channels=self.in_channels,
+            out_channels=self.out_channels,
+        )
+
+    def forward_train(
+        self,
+        pc: Optional[Tensor],
+        ref_kps: Optional[Tensor],
+        ori_kps: Optional[Tensor],
+        aud_feat: Optional[Tensor],
+        mode: str = 'train',
+        return_intermediate_steps: bool = False
+    ):
+
+        # Normalize colors and convert to tensor
+        x_0 = pc
+        B, Nf, Np, D = x_0.shape# batch, nums of frames, nums of points, 3
+
+
+        x_0=x_0[:,:,:,0]# batch, nums of frames, 70
+
+        # Sample random noise
+        noise = torch.randn_like(x_0)
+
+        # Sample random timesteps for each point_cloud
+        timestep = torch.randint(0, self.scheduler.num_train_timesteps, (B,),
+            device=self.device, dtype=torch.long)
+
+        # Add noise to points
+        x_t = self.scheduler.add_noise(x_0, noise, timestep)
+
+        # Conditioning
+        ref_kps = ref_kps[:, :, 0]
+
+        x_t_input = torch.cat([ori_kps.unsqueeze(1), ref_kps.unsqueeze(1), x_t], dim=1)
+
+        # x_t_input = torch.cat([ref_kps.unsqueeze(1), x_t], dim=1)
+
+        # ori_kps_repeat = torch.repeat_interleave(ori_kps.unsqueeze(1), repeats=Nf+1, dim=1)
+
+        # x_t_input = torch.cat([x_t_input, ori_kps_repeat], dim=-1)  #B, 32+1, 51+45
+
+
+        aud_feat = torch.cat([torch.zeros(B, 2, 512).cuda(), aud_feat], 1)
+
+        # Augmentation for audio feature
+        if mode in 'train':
+            if torch.rand(1) > 0.3:
+                mean = torch.mean(aud_feat)
+                std = torch.std(aud_feat)
+                sample = torch.normal(mean=torch.full(aud_feat.shape, mean), std=torch.full(aud_feat.shape, std)).cuda()
+                aud_feat = sample + aud_feat
+            else:
+                pass
+        else:
+            pass
+
+        # Forward
+        noise_pred = self.point_model(x_t_input, timestep, context=aud_feat)    #torch.cat([mel_feat,style_embed],-1))
+        noise_pred = noise_pred[:, 2:]
+        #
+        # Check
+        if not noise_pred.shape == noise.shape:
+            raise ValueError(f'{noise_pred.shape=} and {noise.shape=}')
+
+        # Loss
+        loss = F.mse_loss(noise_pred, noise)
+
+        loss_pose = F.mse_loss(noise_pred[:, :, :6], noise[:, :, :6])
+        loss_exp = F.mse_loss(noise_pred[:, :, 6:], noise[:, :, 6:])
+
+
+        # Whether to return intermediate steps
+        if return_intermediate_steps:
+            return loss, (x_0, x_t, noise, noise_pred)
+
+        return loss, loss_exp, loss_pose
+
+    # def forward_train(
+    #     self,
+    #     pc: Optional[Tensor],
+    #     ref_kps: Optional[Tensor],
+    #     ori_kps: Optional[Tensor],
+    #     aud_feat: Optional[Tensor],
+    #     mode: str = 'train',
+    #     return_intermediate_steps: bool = False
+    # ):
+    #
+    #     # Normalize colors and convert to tensor
+    #     x_0 = pc
+    #     B, Nf, Np, D = x_0.shape# batch, nums of frames, nums of points, 3
+    #
+    #     # ori_kps = torch.repeat_interleave(ori_kps.unsqueeze(1), Nf, dim=1)      # B, Nf, 45
+    #     #
+    #     # ref_kps = ref_kps[:, :, 0]
+    #     # ref_kps = torch.repeat_interleave(ref_kps.unsqueeze(1), Nf, dim=1)  # B, Nf, 91
+    #
+    #     x_0 = x_0[:,:,:,0]
+    #
+    #     # Sample random noise
+    #     noise = torch.randn_like(x_0)
+    #
+    #     # Sample random timesteps for each point_cloud
+    #     timestep = torch.randint(0, self.scheduler.num_train_timesteps, (B,),
+    #         device=self.device, dtype=torch.long)
+    #
+    #     # Add noise to points
+    #     x_t = self.scheduler.add_noise(x_0, noise, timestep)
+    #
+    #     # Conditioning
+    #     ref_kps = ref_kps[:,:,0]
+    #
+    #     # x_t_input = torch.cat([ref_kps.unsqueeze(1), x_t], dim=1)
+    #
+    #     # x_0 = torch.cat([x_0, ref_kps, ori_kps], dim=2)  # B, Nf, 91+91+45
+    #
+    #     x_t_input = torch.cat([ref_kps.unsqueeze(1), x_t], dim=1)
+    #     # x_t_input = torch.cat([ori_kps.unsqueeze(1), ref_kps.unsqueeze(1), x_t], dim=1)
+    #
+    #     aud_feat = torch.cat([torch.zeros(B, 1, 512).cuda(), aud_feat], 1)
+    #
+    #     # Augmentation for audio feature
+    #     if mode in 'train':
+    #         if torch.rand(1) > 0.3:
+    #             mean = torch.mean(aud_feat)
+    #             std = torch.std(aud_feat)
+    #             sample = torch.normal(mean=torch.full(aud_feat.shape, mean), std=torch.full(aud_feat.shape, std)).cuda()
+    #             aud_feat = sample + aud_feat
+    #         else:
+    #             pass
+    #     else:
+    #         pass
+    #
+    #     # Forward
+    #     noise_pred = self.point_model(x_t_input, timestep, context=aud_feat)
+    #     noise_pred = noise_pred[:, 1:]
+    #
+    #     # Check
+    #     # if not noise_pred.shape == noise.shape:
+    #     #     raise ValueError(f'{noise_pred.shape=} and {noise.shape=}')
+    #
+    #     # Loss
+    #     loss = F.mse_loss(noise_pred, noise)
+    #
+    #     # loss_kp = F.mse_loss(noise_pred[:, :, :45], noise[:, :, :45])
+    #
+    #     # Whether to return intermediate steps
+    #     if return_intermediate_steps:
+    #         return loss, (x_0, x_t, noise, noise_pred)
+    #
+    #     return loss
+
+    # @torch.no_grad()
+    # def forward_sample(
+    #         self,
+    #         num_points: int,
+    #         ref_kps: Optional[Tensor],
+    #         ori_kps: Optional[Tensor],
+    #         aud_feat: Optional[Tensor],
+    #         # Optional overrides
+    #         scheduler: Optional[str] = 'ddpm',
+    #         # Inference parameters
+    #         num_inference_steps: Optional[int] = 1000,
+    #         eta: Optional[float] = 0.0,  # for DDIM
+    #         # Whether to return all the intermediate steps in generation
+    #         return_sample_every_n_steps: int = -1,
+    #         # Whether to disable tqdm
+    #         disable_tqdm: bool = False,
+    # ):
+    #
+    #     # Get scheduler from mapping, or use self.scheduler if None
+    #     scheduler = self.scheduler if scheduler is None else self.schedulers_map[scheduler]
+    #
+    #     # Get the size of the noise
+    #     Np = num_points
+    #     Nf = aud_feat.size(1)
+    #     B = 1
+    #     D = 3
+    #     device = self.device
+    #
+    #     # Sample noise
+    #     x_t = torch.randn(B, Nf, Np, D, device=device)
+    #
+    #     x_t = x_t[:, :, :, 0]
+    #
+    #     # ori_kps = torch.repeat_interleave(ori_kps.unsqueeze(1), Nf, dim=1)  # B, Nf, 45
+    #
+    #     ref_kps = ref_kps[:, :, 0]
+    #     # ref_kps = torch.repeat_interleave(ref_kps.unsqueeze(1), Nf, dim=1)  # B, Nf, 91
+    #
+    #     # Set timesteps
+    #     accepts_offset = "offset" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+    #     extra_set_kwargs = {"offset": 1} if accepts_offset else {}
+    #     scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)
+    #
+    #     # Prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+    #     # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+    #     # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+    #     # and should be between [0, 1]
+    #     accepts_eta = "eta" in set(inspect.signature(scheduler.step).parameters.keys())
+    #     extra_step_kwargs = {"eta": eta} if accepts_eta else {}
+    #
+    #     # Loop over timesteps
+    #     all_outputs = []
+    #     return_all_outputs = (return_sample_every_n_steps > 0)
+    #     progress_bar = tqdm(scheduler.timesteps.to(device), desc=f'Sampling ({x_t.shape})', disable=disable_tqdm)
+    #
+    #     # ori_kps = torch.repeat_interleave(ori_kps[:, 6:].unsqueeze(1), Nf + 1, dim=1)
+    #     aud_feat = torch.cat([torch.zeros(B, 1, 512).cuda(), aud_feat], 1)
+    #     # aud_feat = torch.cat([ori_kps, aud_feat], -1)
+    #
+    #     # aud_feat = torch.cat([torch.zeros(B, 2, 512).cuda(), aud_feat], 1)
+    #
+    #     for i, t in enumerate(progress_bar):
+    #
+    #         # Conditioning
+    #         x_t_input = torch.cat([ref_kps.unsqueeze(1).detach(), x_t], dim=1)
+    #         # x_t_input = torch.cat([ori_kps.unsqueeze(1).detach(), ref_kps.unsqueeze(1).detach(), x_t], dim=1)
+    #         # x_t_input = torch.cat([x_t, ref_kps, ori_kps], dim=2)           # B, Nf, 91+91+45
+    #
+    #         # Forward
+    #         # noise_pred = self.point_model(x_t_input, t.reshape(1).expand(B), context=aud_feat)[:, 1:]
+    #         noise_pred = self.point_model(x_t_input, t.reshape(1).expand(B), context=aud_feat)[:, 1:]
+    #
+    #         # noise_pred = noise_pred[:, :, :51]
+    #
+    #         # Step
+    #         # x_t = x_t[:, :, :51]
+    #         x_t = scheduler.step(noise_pred, t, x_t, **extra_step_kwargs).prev_sample
+    #
+    #         # Append to output list if desired
+    #         if (return_all_outputs and (i % return_sample_every_n_steps == 0 or i == len(scheduler.timesteps) - 1)):
+    #             all_outputs.append(x_t)
+    #
+    #     # Convert output back into a point cloud, undoing normalization and scaling
+    #     output = x_t
+    #     output = torch.stack([output, output, output], -1)
+    #     if return_all_outputs:
+    #         all_outputs = torch.stack(all_outputs, dim=1)  # (B, sample_steps, N, D)
+    #     return (output, all_outputs) if return_all_outputs else output
+
+
+    @torch.no_grad()
+    def forward_sample(
+        self,
+        num_points: int,
+        ref_kps: Optional[Tensor],
+        ori_kps: Optional[Tensor],
+        aud_feat: Optional[Tensor],
+        # Optional overrides
+        scheduler: Optional[str] = 'ddpm',
+        # Inference parameters
+        num_inference_steps: Optional[int] = 1000,
+        eta: Optional[float] = 0.0,  # for DDIM
+        # Whether to return all the intermediate steps in generation
+        return_sample_every_n_steps: int = -1,
+        # Whether to disable tqdm
+        disable_tqdm: bool = False,
+    ):
+
+        # Get scheduler from mapping, or use self.scheduler if None
+        scheduler = self.scheduler if scheduler is None else self.schedulers_map[scheduler]
+
+        # Get the size of the noise
+        Np = num_points
+        Nf = aud_feat.size(1)
+        B = 1
+        D = 3
+        device = self.device
+
+        # Sample noise
+        x_t = torch.randn(B, Nf, Np, D, device=device)
+
+        x_t = x_t[:, :, :, 0]
+
+        ref_kps = ref_kps[:,:,0]
+
+        # Set timesteps
+        accepts_offset = "offset" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        extra_set_kwargs = {"offset": 1} if accepts_offset else {}
+        scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)
+
+        # Prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+        accepts_eta = "eta" in set(inspect.signature(scheduler.step).parameters.keys())
+        extra_step_kwargs = {"eta": eta} if accepts_eta else {}
+
+        # Loop over timesteps
+        all_outputs = []
+        return_all_outputs = (return_sample_every_n_steps > 0)
+        progress_bar = tqdm(scheduler.timesteps.to(device), desc=f'Sampling ({x_t.shape})', disable=disable_tqdm)
+
+        # ori_kps = torch.repeat_interleave(ori_kps[:, 6:].unsqueeze(1), Nf + 1, dim=1)
+        # aud_feat = torch.cat([torch.zeros(B, 1, 512).cuda(), aud_feat], 1)
+        # aud_feat = torch.cat([ori_kps, aud_feat], -1)
+
+        aud_feat = torch.cat([torch.zeros(B, 2, 512).cuda(), aud_feat], 1)
+
+        for i, t in enumerate(progress_bar):
+
+            # Conditioning
+            # x_t_input = torch.cat([ref_kps.unsqueeze(1), x_t], dim=1)
+            #
+            # ori_kps_repeat = torch.repeat_interleave(ori_kps.unsqueeze(1), repeats=Nf + 1, dim=1)
+            #
+            # x_t_input = torch.cat([x_t_input.detach(), ori_kps_repeat.detach()], dim=-1)  # B, 32+1, 51+45
+
+
+            x_t_input = torch.cat([ori_kps.unsqueeze(1).detach(),ref_kps.unsqueeze(1).detach(), x_t], dim=1)
+            # x_t_input = torch.cat([ref_kps.unsqueeze(1).detach(), x_t], dim=1)
+
+
+            # Forward
+            # noise_pred = self.point_model(x_t_input, t.reshape(1).expand(B), context=aud_feat)[:, 1:]
+            noise_pred = self.point_model(x_t_input, t.reshape(1).expand(B), context=aud_feat)[:, 2:]
+
+
+            # Step
+            x_t = scheduler.step(noise_pred, t, x_t, **extra_step_kwargs).prev_sample
+
+            # Append to output list if desired
+            if (return_all_outputs and (i % return_sample_every_n_steps == 0 or i == len(scheduler.timesteps) - 1)):
+                all_outputs.append(x_t)
+
+        # Convert output back into a point cloud, undoing normalization and scaling
+        output = x_t
+        output = torch.stack([output,output,output],-1)
+        if return_all_outputs:
+            all_outputs = torch.stack(all_outputs, dim=1)  # (B, sample_steps, N, D)
+        return (output, all_outputs) if return_all_outputs else output
+
+    def forward(self, batch: dict, mode: str = 'train', **kwargs):
+        """A wrapper around the forward method for training and inference"""
+
+        if mode == 'train':
+            return self.forward_train(
+                pc=batch['sequence_keypoints'],
+                ref_kps=batch['ref_keypoint'],
+                ori_kps=batch['ori_keypoint'],
+                aud_feat=batch['aud_feat'],
+                mode='train',
+                **kwargs)
+        elif mode == 'val':
+            return self.forward_train(
+                pc=batch['sequence_keypoints'],
+                ref_kps=batch['ref_keypoint'],
+                ori_kps=batch['ori_keypoint'],
+                aud_feat=batch['aud_feat'],
+                mode='val',
+                **kwargs)
+        elif mode == 'sample':
+            num_points = 68
+            return self.forward_sample(
+                num_points=num_points,
+                ref_kps=batch['ref_keypoint'],
+                ori_kps=batch['ori_keypoint'],
+                aud_feat=batch['aud_feat'],
+                **kwargs) 
+        else:
+            raise NotImplementedError()

difpoint/model/model_utils.py

@@ -0,0 +1,35 @@
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+
+def set_requires_grad(module: nn.Module, requires_grad: bool):
+    for p in module.parameters():
+        p.requires_grad_(requires_grad)
+
+
+def compute_distance_transform(mask: torch.Tensor):
+    image_size = mask.shape[-1]
+    distance_transform = torch.stack([
+        torch.from_numpy(cv2.distanceTransform(
+            (1 - m), distanceType=cv2.DIST_L2, maskSize=cv2.DIST_MASK_3
+        ) / (image_size / 2))
+        for m in mask.squeeze(1).detach().cpu().numpy().astype(np.uint8)
+    ]).unsqueeze(1).clip(0, 1).to(mask.device)
+    return distance_transform
+
+
+def default(x, d):
+    return d if x is None else x
+
+def get_custom_betas(beta_start: float, beta_end: float, warmup_frac: float = 0.3, num_train_timesteps: int = 1000):
+    """Custom beta schedule"""
+    betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32)
+    warmup_frac = 0.3
+    warmup_time = int(num_train_timesteps * warmup_frac)
+    warmup_steps = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64)
+    warmup_time = min(warmup_time, num_train_timesteps)
+    betas[:warmup_time] = warmup_steps[:warmup_time]
+    return betas

difpoint/model/point_model.py

@@ -0,0 +1,38 @@
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers import ModelMixin
+from torch import Tensor
+
+from .temporaltrans.temptrans import SimpleTemperalPointModel, SimpleTransModel
+
+class PointModel(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        model_type: str = 'pvcnn',
+        in_channels: int = 3,
+        out_channels: int = 3,
+        embed_dim: int = 64,
+        dropout: float = 0.1,
+        width_multiplier: int = 1,
+        voxel_resolution_multiplier: int = 1,
+    ):
+        super().__init__()
+        self.model_type = model_type
+        if self.model_type == 'simple':
+            self.autocast_context = torch.autocast('cuda', dtype=torch.float32)
+            self.model = SimpleTransModel(
+                embed_dim=embed_dim,
+                num_classes=out_channels,
+                extra_feature_channels=(in_channels - 3),
+            )
+            self.model.output_projection.bias.data.normal_(0, 1e-6)
+            self.model.output_projection.weight.data.normal_(0, 1e-6)
+        else:
+            raise NotImplementedError()
+
+    def forward(self, inputs: Tensor, t: Tensor, context=None) -> Tensor:
+        """ Receives input of shape (B, N, in_channels) and returns output
+            of shape (B, N, out_channels) """
+        with self.autocast_context:
+            return self.model(inputs, t, context)

difpoint/model/temporaltrans/pointnet_util.py

@@ -0,0 +1,311 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from time import time
+import numpy as np
+
+
+# reference https://github.com/yanx27/Pointnet_Pointnet2_pytorch, modified by Yang You
+
+
+def timeit(tag, t):
+    print("{}: {}s".format(tag, time() - t))
+    return time()
+
+def pc_normalize(pc):
+    centroid = np.mean(pc, axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
+    pc = pc / m
+    return pc
+
+def square_distance(src, dst):
+    """
+    Calculate Euclid distance between each two points.
+    src^T * dst = xn * xm + yn * ym + zn * zm；
+    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
+    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
+    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
+         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
+    Input:
+        src: source points, [B, N, C]
+        dst: target points, [B, M, C]
+    Output:
+        dist: per-point square distance, [B, N, M]
+    """
+    return torch.sum((src[:, :, None] - dst[:, None]) ** 2, dim=-1)
+
+
+def index_points(points, idx):
+    """
+    Input:
+        points: input points data, [B, N, C]
+        idx: sample index data, [B, S, [K]]
+    Return:
+        new_points:, indexed points data, [B, S, [K], C]
+    """
+    raw_size = idx.size()
+    idx = idx.reshape(raw_size[0], -1)
+    res = torch.gather(points, 1, idx[..., None].expand(-1, -1, points.size(-1)))
+    return res.reshape(*raw_size, -1)
+
+
+def farthest_point_sample(xyz, npoint):
+    """
+    Input:
+        xyz: pointcloud data, [B, N, 3]
+        npoint: number of samples
+    Return:
+        centroids: sampled pointcloud index, [B, npoint]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
+    distance = torch.ones(B, N).to(device) * 1e10
+    farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
+    batch_indices = torch.arange(B, dtype=torch.long).to(device)
+    for i in range(npoint):
+        centroids[:, i] = farthest
+        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
+        dist = torch.sum((xyz - centroid) ** 2, -1)
+        distance = torch.min(distance, dist)
+        farthest = torch.max(distance, -1)[1]
+    return centroids
+
+
+def query_ball_point(radius, nsample, xyz, new_xyz):
+    """
+    Input:
+        radius: local region radius
+        nsample: max sample number in local region
+        xyz: all points, [B, N, 3]
+        new_xyz: query points, [B, S, 3]
+    Return:
+        group_idx: grouped points index, [B, S, nsample]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    _, S, _ = new_xyz.shape
+    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
+    sqrdists = square_distance(new_xyz, xyz)
+    group_idx[sqrdists > radius ** 2] = N
+    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
+    group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
+    mask = group_idx == N
+    group_idx[mask] = group_first[mask]
+    return group_idx
+
+
+def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False, knn=False):
+    """
+    Input:
+        npoint:
+        radius:
+        nsample:
+        xyz: input points position data, [B, N, 3]
+        points: input points data, [B, N, D]
+    Return:
+        new_xyz: sampled points position data, [B, npoint, nsample, 3]
+        new_points: sampled points data, [B, npoint, nsample, 3+D]
+    """
+    B, N, C = xyz.shape
+    S = npoint
+    fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint]
+    torch.cuda.empty_cache()
+    new_xyz = index_points(xyz, fps_idx)
+    torch.cuda.empty_cache()
+    if knn:
+        dists = square_distance(new_xyz, xyz)  # B x npoint x N
+        idx = dists.argsort()[:, :, :nsample]  # B x npoint x K
+    else:
+        idx = query_ball_point(radius, nsample, xyz, new_xyz)
+    torch.cuda.empty_cache()
+    grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
+    torch.cuda.empty_cache()
+    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
+    torch.cuda.empty_cache()
+
+    if points is not None:
+        grouped_points = index_points(points, idx)
+        new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
+    else:
+        new_points = grouped_xyz_norm
+    if returnfps:
+        return new_xyz, new_points, grouped_xyz, fps_idx
+    else:
+        return new_xyz, new_points
+
+
+def sample_and_group_all(xyz, points):
+    """
+    Input:
+        xyz: input points position data, [B, N, 3]
+        points: input points data, [B, N, D]
+    Return:
+        new_xyz: sampled points position data, [B, 1, 3]
+        new_points: sampled points data, [B, 1, N, 3+D]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    new_xyz = torch.zeros(B, 1, C).to(device)
+    grouped_xyz = xyz.view(B, 1, N, C)
+    if points is not None:
+        new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
+    else:
+        new_points = grouped_xyz
+    return new_xyz, new_points
+
+
+class PointNetSetAbstraction(nn.Module):
+    def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all, knn=False):
+        super(PointNetSetAbstraction, self).__init__()
+        self.npoint = npoint
+        self.radius = radius
+        self.nsample = nsample
+        self.knn = knn
+        self.mlp_convs = nn.ModuleList()
+        self.mlp_bns = nn.ModuleList()
+        last_channel = in_channel
+        for out_channel in mlp:
+            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
+            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
+            last_channel = out_channel
+        self.group_all = group_all
+
+    def forward(self, xyz, points):
+        """
+        Input:
+            xyz: input points position data, [B, N, C]
+            points: input points data, [B, N, C]
+        Return:
+            new_xyz: sampled points position data, [B, S, C]
+            new_points_concat: sample points feature data, [B, S, D']
+        """
+        if self.group_all:
+            new_xyz, new_points = sample_and_group_all(xyz, points)
+        else:
+            new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points, knn=self.knn)
+        # new_xyz: sampled points position data, [B, npoint, C]
+        # new_points: sampled points data, [B, npoint, nsample, C+D]
+        new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
+        for i, conv in enumerate(self.mlp_convs):
+            bn = self.mlp_bns[i]
+            new_points =  F.relu(bn(conv(new_points)))
+
+        new_points = torch.max(new_points, 2)[0].transpose(1, 2)
+        return new_xyz, new_points
+
+
+class PointNetSetAbstractionMsg(nn.Module):
+    def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list, knn=False):
+        super(PointNetSetAbstractionMsg, self).__init__()
+        self.npoint = npoint
+        self.radius_list = radius_list
+        self.nsample_list = nsample_list
+        self.knn = knn
+        self.conv_blocks = nn.ModuleList()
+        self.bn_blocks = nn.ModuleList()
+        for i in range(len(mlp_list)):
+            convs = nn.ModuleList()
+            bns = nn.ModuleList()
+            last_channel = in_channel + 3
+            for out_channel in mlp_list[i]:
+                convs.append(nn.Conv2d(last_channel, out_channel, 1))
+                bns.append(nn.BatchNorm2d(out_channel))
+                last_channel = out_channel
+            self.conv_blocks.append(convs)
+            self.bn_blocks.append(bns)
+
+    def forward(self, xyz, points, seed_idx=None):
+        """
+        Input:
+            xyz: input points position data, [B, C, N]
+            points: input points data, [B, D, N]
+        Return:
+            new_xyz: sampled points position data, [B, C, S]
+            new_points_concat: sample points feature data, [B, D', S]
+        """
+
+        B, N, C = xyz.shape
+        S = self.npoint
+        new_xyz = index_points(xyz, farthest_point_sample(xyz, S) if seed_idx is None else seed_idx)
+        new_points_list = []
+        for i, radius in enumerate(self.radius_list):
+            K = self.nsample_list[i]
+            if self.knn:
+                dists = square_distance(new_xyz, xyz)  # B x npoint x N
+                group_idx = dists.argsort()[:, :, :K]  # B x npoint x K
+            else:
+                group_idx = query_ball_point(radius, K, xyz, new_xyz)
+            grouped_xyz = index_points(xyz, group_idx)
+            grouped_xyz -= new_xyz.view(B, S, 1, C)
+            if points is not None:
+                grouped_points = index_points(points, group_idx)
+                grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
+            else:
+                grouped_points = grouped_xyz
+
+            grouped_points = grouped_points.permute(0, 3, 2, 1)  # [B, D, K, S]
+            for j in range(len(self.conv_blocks[i])):
+                conv = self.conv_blocks[i][j]
+                bn = self.bn_blocks[i][j]
+                grouped_points =  F.relu(bn(conv(grouped_points)))
+            new_points = torch.max(grouped_points, 2)[0]  # [B, D', S]
+            new_points_list.append(new_points)
+
+        new_points_concat = torch.cat(new_points_list, dim=1).transpose(1, 2)
+        return new_xyz, new_points_concat
+
+
+# NoteL this function swaps N and C
+class PointNetFeaturePropagation(nn.Module):
+    def __init__(self, in_channel, mlp):
+        super(PointNetFeaturePropagation, self).__init__()
+        self.mlp_convs = nn.ModuleList()
+        self.mlp_bns = nn.ModuleList()
+        last_channel = in_channel
+        for out_channel in mlp:
+            self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
+            self.mlp_bns.append(nn.BatchNorm1d(out_channel))
+            last_channel = out_channel
+
+    def forward(self, xyz1, xyz2, points1, points2):
+        """
+        Input:
+            xyz1: input points position data, [B, C, N]
+            xyz2: sampled input points position data, [B, C, S]
+            points1: input points data, [B, D, N]
+            points2: input points data, [B, D, S]
+        Return:
+            new_points: upsampled points data, [B, D', N]
+        """
+        xyz1 = xyz1.permute(0, 2, 1)
+        xyz2 = xyz2.permute(0, 2, 1)
+
+        points2 = points2.permute(0, 2, 1)
+        B, N, C = xyz1.shape
+        _, S, _ = xyz2.shape
+
+        if S == 1:
+            interpolated_points = points2.repeat(1, N, 1)
+        else:
+            dists = square_distance(xyz1, xyz2)
+            dists, idx = dists.sort(dim=-1)
+            dists, idx = dists[:, :, :3], idx[:, :, :3]  # [B, N, 3]
+
+            dist_recip = 1.0 / (dists + 1e-8)
+            norm = torch.sum(dist_recip, dim=2, keepdim=True)
+            weight = dist_recip / norm
+            interpolated_points = torch.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2)
+
+        if points1 is not None:
+            points1 = points1.permute(0, 2, 1)
+            new_points = torch.cat([points1, interpolated_points], dim=-1)
+        else:
+            new_points = interpolated_points
+
+        new_points = new_points.permute(0, 2, 1)
+        for i, conv in enumerate(self.mlp_convs):
+            bn = self.mlp_bns[i]
+            new_points = F.relu(bn(conv(new_points)))
+        return new_points

difpoint/model/temporaltrans/pointtransformerv2.py

@@ -0,0 +1,250 @@
+from .transformer_utils import BaseTemperalPointModel
+from copy import deepcopy
+import torch
+import einops
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange
+import pointops
+from pointcept.models.utils import offset2batch, batch2offset
+class PointBatchNorm(nn.Module):
+    """
+    Batch Normalization for Point Clouds data in shape of [B*N, C], [B*N, L, C]
+    """
+
+    def __init__(self, embed_channels):
+        super().__init__()
+        self.norm = nn.BatchNorm1d(embed_channels)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if input.dim() == 3:
+            return (
+                self.norm(input.transpose(1, 2).contiguous())
+                .transpose(1, 2)
+                .contiguous()
+            )
+        elif input.dim() == 2:
+            return self.norm(input)
+        else:
+            raise NotImplementedError
+#https://github.com/Pointcept/Pointcept/blob/main/pointcept/models/point_transformer_v2/point_transformer_v2m2_base.py
+class GroupedVectorAttention(nn.Module):
+    def __init__(
+        self,
+        embed_channels,
+        groups,
+        attn_drop_rate=0.0,
+        qkv_bias=True,
+        pe_multiplier=False,
+        pe_bias=True,
+    ):
+        super(GroupedVectorAttention, self).__init__()
+        self.embed_channels = embed_channels
+        self.groups = groups
+        assert embed_channels % groups == 0
+        self.attn_drop_rate = attn_drop_rate
+        self.qkv_bias = qkv_bias
+        self.pe_multiplier = pe_multiplier
+        self.pe_bias = pe_bias
+
+        self.linear_q = nn.Sequential(
+            nn.Linear(embed_channels, embed_channels, bias=qkv_bias),
+            PointBatchNorm(embed_channels),
+            nn.ReLU(inplace=True),
+        )
+        self.linear_k = nn.Sequential(
+            nn.Linear(embed_channels, embed_channels, bias=qkv_bias),
+            PointBatchNorm(embed_channels),
+            nn.ReLU(inplace=True),
+        )
+
+        self.linear_v = nn.Linear(embed_channels, embed_channels, bias=qkv_bias)
+
+        if self.pe_multiplier:
+            self.linear_p_multiplier = nn.Sequential(
+                nn.Linear(3, embed_channels),
+                PointBatchNorm(embed_channels),
+                nn.ReLU(inplace=True),
+                nn.Linear(embed_channels, embed_channels),
+            )
+        if self.pe_bias:
+            self.linear_p_bias = nn.Sequential(
+                nn.Linear(3, embed_channels),
+                PointBatchNorm(embed_channels),
+                nn.ReLU(inplace=True),
+                nn.Linear(embed_channels, embed_channels),
+            )
+        self.weight_encoding = nn.Sequential(
+            nn.Linear(embed_channels, groups),
+            PointBatchNorm(groups),
+            nn.ReLU(inplace=True),
+            nn.Linear(groups, groups),
+        )
+        self.softmax = nn.Softmax(dim=1)
+        self.attn_drop = nn.Dropout(attn_drop_rate)
+
+    def forward(self, feat, coord, reference_index):
+        query, key, value = (
+            self.linear_q(feat),
+            self.linear_k(feat),
+            self.linear_v(feat),
+        )
+        key = pointops.grouping(reference_index, key, coord, with_xyz=True)
+        value = pointops.grouping(reference_index, value, coord, with_xyz=False)
+        pos, key = key[:, :, 0:3], key[:, :, 3:]
+        relation_qk = key - query.unsqueeze(1)
+        if self.pe_multiplier:
+            pem = self.linear_p_multiplier(pos)
+            relation_qk = relation_qk * pem
+        if self.pe_bias:
+            peb = self.linear_p_bias(pos)
+            relation_qk = relation_qk + peb
+            value = value + peb
+
+        weight = self.weight_encoding(relation_qk)
+        weight = self.attn_drop(self.softmax(weight))
+
+        mask = torch.sign(reference_index + 1)
+        weight = torch.einsum("n s g, n s -> n s g", weight, mask)
+        value = einops.rearrange(value, "n ns (g i) -> n ns g i", g=self.groups)
+        feat = torch.einsum("n s g i, n s g -> n g i", value, weight)
+        feat = einops.rearrange(feat, "n g i -> n (g i)")
+        return feat
+
+class BlockSequence(nn.Module):
+    def __init__(
+        self,
+        depth,
+        embed_channels,
+        groups,
+        neighbours=16,
+        qkv_bias=True,
+        pe_multiplier=False,
+        pe_bias=True,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        enable_checkpoint=False,
+    ):
+        super(BlockSequence, self).__init__()
+
+        if isinstance(drop_path_rate, list):
+            drop_path_rates = drop_path_rate
+            assert len(drop_path_rates) == depth
+        elif isinstance(drop_path_rate, float):
+            drop_path_rates = [deepcopy(drop_path_rate) for _ in range(depth)]
+        else:
+            drop_path_rates = [0.0 for _ in range(depth)]
+
+        self.neighbours = neighbours
+        self.blocks = nn.ModuleList()
+        for i in range(depth):
+            block = Block(
+                embed_channels=embed_channels,
+                groups=groups,
+                qkv_bias=qkv_bias,
+                pe_multiplier=pe_multiplier,
+                pe_bias=pe_bias,
+                attn_drop_rate=attn_drop_rate,
+                drop_path_rate=drop_path_rates[i],
+                enable_checkpoint=enable_checkpoint,
+            )
+            self.blocks.append(block)
+
+    def forward(self, points):
+        coord, feat, offset = points
+        # reference index query of neighbourhood attention
+        # for windows attention, modify reference index query method
+        reference_index, _ = pointops.knn_query(self.neighbours, coord, offset)
+        for block in self.blocks:
+            points = block(points, reference_index)
+        return points
+
+class GVAPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        depth,
+        in_channels,
+        embed_channels,
+        groups,
+        neighbours=16,
+        qkv_bias=True,
+        pe_multiplier=False,
+        pe_bias=True,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        enable_checkpoint=False,
+    ):
+        super(GVAPatchEmbed, self).__init__()
+        self.in_channels = in_channels
+        self.embed_channels = embed_channels
+        self.proj = nn.Sequential(
+            nn.Linear(in_channels, embed_channels, bias=False),
+            PointBatchNorm(embed_channels),
+            nn.ReLU(inplace=True),
+        )
+        self.blocks = BlockSequence(
+            depth=depth,
+            embed_channels=embed_channels,
+            groups=groups,
+            neighbours=neighbours,
+            qkv_bias=qkv_bias,
+            pe_multiplier=pe_multiplier,
+            pe_bias=pe_bias,
+            attn_drop_rate=attn_drop_rate,
+            drop_path_rate=drop_path_rate,
+            enable_checkpoint=enable_checkpoint,
+        )
+
+    def forward(self, points):
+        coord, feat, offset = points
+        feat = self.proj(feat)
+        return self.blocks([coord, feat, offset])
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        embed_channels,
+        groups,
+        qkv_bias=True,
+        pe_multiplier=False,
+        pe_bias=True,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        enable_checkpoint=False,
+    ):
+        super(Block, self).__init__()
+        self.attn = GroupedVectorAttention(
+            embed_channels=embed_channels,
+            groups=groups,
+            qkv_bias=qkv_bias,
+            attn_drop_rate=attn_drop_rate,
+            pe_multiplier=pe_multiplier,
+            pe_bias=pe_bias,
+        )
+        self.fc1 = nn.Linear(embed_channels, embed_channels, bias=False)
+        self.fc3 = nn.Linear(embed_channels, embed_channels, bias=False)
+        self.norm1 = PointBatchNorm(embed_channels)
+        self.norm2 = PointBatchNorm(embed_channels)
+        self.norm3 = PointBatchNorm(embed_channels)
+        self.act = nn.ReLU(inplace=True)
+        self.enable_checkpoint = enable_checkpoint
+        self.drop_path = (
+            DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+        )
+
+    def forward(self, points, reference_index):
+        coord, feat, offset = points
+        identity = feat
+        feat = self.act(self.norm1(self.fc1(feat)))
+        feat = (
+            self.attn(feat, coord, reference_index)
+            if not self.enable_checkpoint
+            else checkpoint(self.attn, feat, coord, reference_index)
+        )
+        feat = self.act(self.norm2(feat))
+        feat = self.norm3(self.fc3(feat))
+        feat = identity + self.drop_path(feat)
+        feat = self.act(feat)
+        return [coord, feat, offset]

difpoint/model/temporaltrans/temptrans.py

@@ -0,0 +1,347 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange
+from .transformer_utils import BaseTemperalPointModel
+import math
+from einops_exts import check_shape, rearrange_many
+from functools import partial
+
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = x[:, None] * emb[None, :]
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+class RelativePositionBias(nn.Module):
+    def __init__(
+        self,
+        heads = 8,
+        num_buckets = 32,
+        max_distance = 128
+    ):
+        super().__init__()
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.relative_attention_bias = nn.Embedding(num_buckets, heads)
+
+    @staticmethod
+    def _relative_position_bucket(relative_position, num_buckets = 32, max_distance = 128):
+        ret = 0
+        n = -relative_position
+
+        num_buckets //= 2
+        ret += (n < 0).long() * num_buckets
+        n = torch.abs(n)
+
+        max_exact = num_buckets // 2
+        is_small = n < max_exact
+
+        val_if_large = max_exact + (
+            torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact)
+        ).long()
+        val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1))
+
+        ret += torch.where(is_small, n, val_if_large)
+        return ret
+
+    def forward(self, n, device):
+        q_pos = torch.arange(n, dtype = torch.long, device = device)
+        k_pos = torch.arange(n, dtype = torch.long, device = device)
+        rel_pos = rearrange(k_pos, 'j -> 1 j') - rearrange(q_pos, 'i -> i 1')
+        rp_bucket = self._relative_position_bucket(rel_pos, num_buckets = self.num_buckets, max_distance = self.max_distance)
+        values = self.relative_attention_bias(rp_bucket)
+        return rearrange(values, 'i j h -> h i j')
+def exists(x):
+    return x is not None
+
+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, eps = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.gamma = nn.Parameter(torch.ones(1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, dim))
+
+    def forward(self, x):
+        var = torch.var(x, dim = -1, unbiased = False, keepdim = True)
+        mean = torch.mean(x, dim = -1, keepdim = True)
+        return (x - mean) / (var + self.eps).sqrt() * self.gamma + self.beta
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = LayerNorm(dim)
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, **kwargs)
+
+
+class EinopsToAndFrom(nn.Module):
+    def __init__(self, from_einops, to_einops, fn):
+        super().__init__()
+        self.from_einops = from_einops
+        self.to_einops = to_einops
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        shape = x.shape
+        reconstitute_kwargs = dict(tuple(zip(self.from_einops.split(' '), shape)))
+        x = rearrange(x, f'{self.from_einops} -> {self.to_einops}')
+        x = self.fn(x, **kwargs)
+        x = rearrange(x, f'{self.to_einops} -> {self.from_einops}', **reconstitute_kwargs)
+        return x
+
+class Attention(nn.Module):
+    def __init__(
+            self, dim, heads=4, attn_head_dim=None, casual_attn=False,rotary_emb = None):
+        super().__init__()
+        self.num_heads = heads
+        head_dim = dim // heads
+        self.casual_attn = casual_attn
+
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+
+        all_head_dim = head_dim * self.num_heads
+        self.scale = head_dim ** -0.5
+        self.to_qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+        self.proj = nn.Linear(all_head_dim, dim)
+        self.rotary_emb = rotary_emb
+
+    def forward(self, x, pos_bias = None):
+        N, device = x.shape[-2], x.device
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+
+        q, k, v = rearrange_many(qkv, '... n (h d) -> ... h n d', h=self.num_heads)
+
+        q = q * self.scale
+
+        if exists(self.rotary_emb):
+            q = self.rotary_emb.rotate_queries_or_keys(q)
+            k = self.rotary_emb.rotate_queries_or_keys(k)
+
+        sim = torch.einsum('... h i d, ... h j d -> ... h i j', q, k)
+
+        if exists(pos_bias):
+            sim = sim + pos_bias
+
+        if self.casual_attn:
+            mask = torch.tril(torch.ones(sim.size(-1), sim.size(-2))).to(device)
+            sim = sim.masked_fill(mask[..., :, :] == 0, float('-inf'))
+
+        attn = sim.softmax(dim = -1)
+        x = torch.einsum('... h i j, ... h j d -> ... h i d', attn, v)
+        x = rearrange(x, '... h n d -> ... n (h d)')
+        x = self.proj(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, dim, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim, dim_out)
+        self.norm = LayerNorm(dim)
+        self.act = nn.SiLU()
+
+    def forward(self, x, scale_shift=None):
+        x = self.proj(x)
+
+        if exists(scale_shift):
+            x = self.norm(x)
+            scale, shift = scale_shift
+            x = x * (scale + 1) + shift
+        return self.act(x)
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, dim_out, cond_dim=None):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(cond_dim, dim_out * 2)
+        ) if exists(cond_dim) else None
+
+        self.block1 = Block(dim, dim_out)
+        self.block2 = Block(dim_out, dim_out)
+
+    def forward(self, x, cond_emb=None):
+        scale_shift = None
+        if exists(self.mlp):
+            assert exists(cond_emb), 'time emb must be passed in'
+            cond_emb = self.mlp(cond_emb)
+            #cond_emb = rearrange(cond_emb, 'b f c -> b f 1 c')
+            scale_shift = cond_emb.chunk(2, dim=-1)
+
+        h = self.block1(x, scale_shift=scale_shift)
+        h = self.block2(h)
+        return h + x
+
+from rotary_embedding_torch import RotaryEmbedding
+
+
+class SimpleTransModel(BaseTemperalPointModel):
+    """
+    A simple model that processes a point cloud by applying a series of MLPs to each point
+    individually, along with some pooled global features.
+    """
+
+    def get_layers(self):
+
+
+        # self.input_projection = nn.Linear(
+        #     in_features=51,
+        #     out_features=self.dim
+        # )
+
+        self.input_projection = nn.Linear(
+            in_features=70,
+            out_features=self.dim
+        )
+
+        cond_dim = 512 + self.timestep_embed_dim
+
+        num_head = self.dim//64
+        rotary_emb = RotaryEmbedding(min(32, num_head))
+
+        self.time_rel_pos_bias = RelativePositionBias(heads=num_head, max_distance=128)  # realistically will not be able to generate that many frames of video... yet
+
+        temporal_casual_attn = lambda dim: Attention(dim, heads=num_head, casual_attn=False,rotary_emb=rotary_emb)
+
+        cond_block= partial(ResnetBlock,cond_dim=cond_dim)
+
+        layers = nn.ModuleList([])
+
+        for _ in range(self.num_layers):
+            layers.append(nn.ModuleList([
+                cond_block(self.dim,self.dim),
+                cond_block(self.dim,self.dim),
+                Residual(PreNorm(self.dim,temporal_casual_attn(self.dim)))
+            ]))
+
+        return layers
+
+    def forward(self, inputs: torch.Tensor, timesteps: torch.Tensor, context=None):
+        """
+         Apply the model to an input batch.
+         :param x: an [N x C x ...] Tensor of inputs.
+         :param timesteps: a 1-D batch of timesteps.
+         :param context: conditioning plugged in via crossattn
+         """
+        # Prepare inputs
+
+        batch, num_frames, channels = inputs.size()
+
+        device = inputs.device
+        #assert channels==3
+
+        # Positional encoding of point coords
+        # inputs=rearrange(inputs,'b f p c->(b f) p c')
+        # pos_emb=self.positional_encoding(inputs)
+        x = self.input_projection(inputs)
+        #x = rearrange(x,'(b f) p c-> b f p c',b=batch)
+
+        t_emb = self.time_mlp(timesteps) if exists(self.time_mlp) else None
+        t_emb = t_emb[:,None,:].expand(-1, num_frames, -1)  # b f c
+        if context is not None:
+            t_emb = torch.cat([t_emb, context],-1)
+
+        time_rel_pos_bias = self.time_rel_pos_bias(num_frames, device=device)
+
+        for block1, block2,  temporal_casual_attn in self.layers:
+            x = block1(x, t_emb)
+            x = block2(x, t_emb)
+            x = temporal_casual_attn(x, pos_bias=time_rel_pos_bias)
+
+        # Project
+        x = self.output_projection(x)
+        return x
+
+
+
+class SimpleTemperalPointModel(BaseTemperalPointModel):
+    """
+    A simple model that processes a point cloud by applying a series of MLPs to each point
+    individually, along with some pooled global features.
+    """
+
+    def get_layers(self):
+        audio_dim = 512
+
+        cond_dim = audio_dim + self.timestep_embed_dim
+
+        num_head = 4
+        rotary_emb = RotaryEmbedding(min(32, num_head))
+        self.time_rel_pos_bias = RelativePositionBias(heads=num_head, max_distance=128)  # realistically will not be able to generate that many frames of video... yet
+
+        temporal_casual_attn = lambda dim: EinopsToAndFrom('b f p c', 'b p f c', Attention(dim, heads=num_head, casual_attn=False, rotary_emb = rotary_emb))
+
+        spatial_kp_attn= lambda dim: EinopsToAndFrom('b f p c', 'b f p c', Attention(dim, heads=num_head))
+
+        cond_block= partial(ResnetBlock,cond_dim=cond_dim)
+
+        layers = nn.ModuleList([])
+
+        for _ in range(self.num_layers):
+            layers.append(nn.ModuleList([
+                cond_block(self.dim,self.dim),
+                cond_block(self.dim,self.dim),
+                Residual(PreNorm(self.dim,spatial_kp_attn(self.dim))),
+                Residual(PreNorm(self.dim,temporal_casual_attn(self.dim)))
+            ]))
+
+        return layers
+
+    def forward(self, inputs: torch.Tensor, timesteps: torch.Tensor, context=None):
+        """
+         Apply the model to an input batch.
+         :param x: an [N x C x ...] Tensor of inputs.
+         :param timesteps: a 1-D batch of timesteps.
+         :param context: conditioning plugged in via crossattn
+         """
+        # Prepare inputs
+
+        batch, num_frames, num_points, channels = inputs.size()
+        device = inputs.device
+        #assert channels==3
+
+        # Positional encoding of point coords
+        inputs=rearrange(inputs,'b f p c->(b f) p c')
+        pos_emb=self.positional_encoding(inputs)
+        x = self.input_projection(torch.cat([inputs, pos_emb], -1))
+        x = rearrange(x,'(b f) p c-> b f p c',b=batch)
+
+        t_emb = self.time_mlp(timesteps) if exists(self.time_mlp) else None
+        t_emb = t_emb[:,None,:].expand(-1, num_frames, -1)  # b f c
+        if context is not None:
+            t_emb = torch.cat([t_emb,context],-1)
+
+        time_rel_pos_bias = self.time_rel_pos_bias(num_frames, device=device)
+
+        for block1, block2, spatial_kp_attn, temporal_casual_attn in self.layers:
+            x = block1(x, t_emb)
+            x = block2(x, t_emb)
+            x = spatial_kp_attn(x)
+            x = temporal_casual_attn(x, pos_bias=time_rel_pos_bias)
+
+        # Project
+        x = self.output_projection(x)
+        return x
+

difpoint/model/temporaltrans/transformer_utils.py

@@ -0,0 +1,146 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange
+import math
+from einops_exts import check_shape, rearrange_many
+from torch import Size, Tensor, nn
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = x[:, None] * emb[None, :]
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+def map_positional_encoding(v: Tensor, freq_bands: Tensor) -> Tensor:
+    """Map v to positional encoding representation phi(v)
+
+    Arguments:
+        v (Tensor): input features (B, IFeatures)
+        freq_bands (Tensor): frequency bands (N_freqs, )
+
+    Returns:
+        phi(v) (Tensor): fourrier features (B, 3 + (2 * N_freqs) * 3)
+    """
+    pe = [v]
+    for freq in freq_bands:
+        fv = freq * v
+        pe += [torch.sin(fv), torch.cos(fv)]
+    return torch.cat(pe, dim=-1)
+
+class FeatureMapping(nn.Module):
+    """FeatureMapping nn.Module
+
+    Maps v to features following transformation phi(v)
+
+    Arguments:
+        i_dim (int): input dimensions
+        o_dim (int): output dimensions
+    """
+
+    def __init__(self, i_dim: int, o_dim: int) -> None:
+        super().__init__()
+        self.i_dim = i_dim
+        self.o_dim = o_dim
+
+    def forward(self, v: Tensor) -> Tensor:
+        """FeratureMapping forward pass
+
+        Arguments:
+            v (Tensor): input features (B, IFeatures)
+
+        Returns:
+            phi(v) (Tensor): mapped features (B, OFeatures)
+        """
+        raise NotImplementedError("Forward pass not implemented yet!")
+
+class PositionalEncoding(FeatureMapping):
+    """PositionalEncoding module
+
+    Maps v to positional encoding representation phi(v)
+
+    Arguments:
+        i_dim (int): input dimension for v
+        N_freqs (int): #frequency to sample (default: 10)
+    """
+
+    def __init__(
+        self,
+        i_dim: int,
+        N_freqs: int = 10,
+    ) -> None:
+        super().__init__(i_dim, 3 + (2 * N_freqs) * 3)
+        self.N_freqs = N_freqs
+
+        a, b = 1, self.N_freqs - 1
+        freq_bands = 2 ** torch.linspace(a, b, self.N_freqs)
+        self.register_buffer("freq_bands", freq_bands)
+
+    def forward(self, v: Tensor) -> Tensor:
+        """Map v to positional encoding representation phi(v)
+
+        Arguments:
+            v (Tensor): input features (B, IFeatures)
+
+        Returns:
+            phi(v) (Tensor): fourrier features (B, 3 + (2 * N_freqs) * 3)
+        """
+        return map_positional_encoding(v, self.freq_bands)
+
+class BaseTemperalPointModel(nn.Module):
+    """ A base class providing useful methods for point cloud processing. """
+
+    def __init__(
+        self,
+        *,
+        num_classes,
+        embed_dim,
+        extra_feature_channels,
+        dim: int = 768,
+        num_layers: int = 6
+    ):
+        super().__init__()
+
+        self.extra_feature_channels = extra_feature_channels
+        self.timestep_embed_dim = 256
+        self.output_dim = num_classes
+        self.dim = dim
+        self.num_layers = num_layers
+
+
+        self.time_mlp = nn.Sequential(
+            SinusoidalPosEmb(dim),
+            nn.Linear(dim, self.timestep_embed_dim ),
+            nn.SiLU(),
+            nn.Linear(self.timestep_embed_dim , self.timestep_embed_dim )
+        )
+
+        self.positional_encoding = PositionalEncoding(i_dim=3, N_freqs=10)
+        positional_encoding_d_out = 3 + (2 * 10) * 3
+
+        # Input projection (point coords, point coord encodings, other features, and timestep embeddings)
+
+        self.input_projection = nn.Linear(
+            in_features=(3 + positional_encoding_d_out),
+            out_features=self.dim
+        )#b f p c
+
+        # Transformer layers
+        self.layers = self.get_layers()
+
+        # Output projection
+        self.output_projection = nn.Linear(self.dim, self.output_dim)
+    def get_layers(self):
+        raise NotImplementedError('This method should be implemented by subclasses')
+
+    def forward(self, inputs: torch.Tensor, t: torch.Tensor):
+        raise NotImplementedError('This method should be implemented by subclasses')

difpoint/src/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+# @Author  : wenshao
+# @Email   : [email protected]
+# @Project : FasterLivePortrait
+# @FileName: __init__.py.py

difpoint/src/config/argument_config.py

@@ -0,0 +1,48 @@
+# coding: utf-8
+
+"""
+All configs for user
+"""
+
+from dataclasses import dataclass
+import tyro
+from typing_extensions import Annotated
+from typing import Optional
+from .base_config import PrintableConfig, make_abs_path
+
+
+@dataclass(repr=False)  # use repr from PrintableConfig
+class ArgumentConfig(PrintableConfig):
+    ########## input arguments ##########
+    source_image: Annotated[str, tyro.conf.arg(aliases=["-s"])] = make_abs_path('../../assets/examples/source/s6.jpg')  # path to the source portrait
+    driving_info:  Annotated[str, tyro.conf.arg(aliases=["-d"])] = make_abs_path('../../assets/examples/driving/d12.mp4')  # path to driving video or template (.pkl format)
+    output_dir: Annotated[str, tyro.conf.arg(aliases=["-o"])] = 'animations/'  # directory to save output video
+
+    ########## inference arguments ##########
+    flag_use_half_precision: bool = False  # whether to use half precision (FP16). If black boxes appear, it might be due to GPU incompatibility; set to False.
+    flag_crop_driving_video: bool = False  # whether to crop the driving video, if the given driving info is a video
+    device_id: int = 0 # gpu device id
+    flag_force_cpu: bool = False # force cpu inference, WIP!
+    flag_lip_zero: bool = False # whether let the lip to close state before animation, only take effect when flag_eye_retargeting and flag_lip_retargeting is False
+    flag_eye_retargeting: bool = False # not recommend to be True, WIP
+    flag_lip_retargeting: bool = False # not recommend to be True, WIP
+    flag_stitching: bool = False  # recommend to True if head movement is small, False if head movement is large
+    flag_relative_motion: bool = False  # whether to use relative motion
+    flag_pasteback: bool = False  # whether to paste-back/stitch the animated face cropping from the face-cropping space to the original image space
+    flag_do_crop: bool = False  # whether to crop the source portrait to the face-cropping space
+    flag_do_rot: bool = False  # whether to conduct the rotation when flag_do_crop is True
+
+    ########## crop arguments ##########
+    scale: float = 2.3 # the ratio of face area is smaller if scale is larger
+    vx_ratio: float = 0  # the ratio to move the face to left or right in cropping space
+    vy_ratio: float = -0.125  # the ratio to move the face to up or down in cropping space
+
+    scale_crop_video: float = 2.2 # scale factor for cropping video
+    vx_ratio_crop_video: float = 0. # adjust y offset
+    vy_ratio_crop_video: float = -0.1  # adjust x offset
+
+    ########## gradio arguments ##########
+    server_port: Annotated[int, tyro.conf.arg(aliases=["-p"])]  = 8890 # port for gradio server
+    share: bool = False # whether to share the server to public
+    server_name: Optional[str] = "127.0.0.1"  # set the local server name, "0.0.0.0" to broadcast all
+    flag_do_torch_compile: bool = False  # whether to use torch.compile to accelerate generation

difpoint/src/config/base_config.py

@@ -0,0 +1,29 @@
+# coding: utf-8
+
+"""
+pretty printing class
+"""
+
+from __future__ import annotations
+import os.path as osp
+from typing import Tuple
+
+
+def make_abs_path(fn):
+    return osp.join(osp.dirname(osp.realpath(__file__)), fn)
+
+
+class PrintableConfig:  # pylint: disable=too-few-public-methods
+    """Printable Config defining str function"""
+
+    def __repr__(self):
+        lines = [self.__class__.__name__ + ":"]
+        for key, val in vars(self).items():
+            if isinstance(val, Tuple):
+                flattened_val = "["
+                for item in val:
+                    flattened_val += str(item) + "\n"
+                flattened_val = flattened_val.rstrip("\n")
+                val = flattened_val + "]"
+            lines += f"{key}: {str(val)}".split("\n")
+        return "\n    ".join(lines)

difpoint/src/config/crop_config.py

@@ -0,0 +1,29 @@
+# coding: utf-8
+
+"""
+parameters used for crop faces
+"""
+
+from dataclasses import dataclass
+
+from .base_config import PrintableConfig
+
+
+@dataclass(repr=False)  # use repr from PrintableConfig
+class CropConfig(PrintableConfig):
+    insightface_root: str = "./downloaded_repo/pretrained_weights/insightface"
+    landmark_ckpt_path: str = "./downloaded_repo/pretrained_weights/liveportrait/landmark.onnx"
+    device_id: int = 0  # gpu device id
+    flag_force_cpu: bool = False  # force cpu inference, WIP
+    ########## source image cropping option ##########
+    dsize: int = 512  # crop size
+    scale: float = 2.0  # scale factor
+    vx_ratio: float = 0  # vx ratio
+    vy_ratio: float = -0.125  # vy ratio +up, -down
+    max_face_num: int = 0  # max face number, 0 mean no limit
+
+    ########## driving video auto cropping option ##########
+    scale_crop_video: float = 2.2  # 2.0 # scale factor for cropping video
+    vx_ratio_crop_video: float = 0.0  # adjust y offset
+    vy_ratio_crop_video: float = -0.1  # adjust x offset
+    direction: str = "large-small"  # direction of cropping

difpoint/src/config/inference_config.py

@@ -0,0 +1,52 @@
+# coding: utf-8
+
+"""
+config dataclass used for inference
+"""
+
+import os.path as osp
+import cv2
+from numpy import ndarray
+from dataclasses import dataclass
+from typing import Literal, Tuple
+from .base_config import PrintableConfig, make_abs_path
+
+
+@dataclass(repr=False)  # use repr from PrintableConfig
+class InferenceConfig(PrintableConfig):
+    # MODEL CONFIG, NOT EXPORTED PARAMS
+    models_config: str = make_abs_path('./models.yaml')  # portrait animation config
+    checkpoint_F: str = './downloaded_repo/pretrained_weights/liveportrait/base_models/appearance_feature_extractor.pth'  # path to checkpoint of F
+    checkpoint_M: str = './downloaded_repo/pretrained_weights/liveportrait/base_models/motion_extractor.pth'  # path to checkpoint pf M
+    checkpoint_G: str = './downloaded_repo/pretrained_weights/liveportrait/base_models/spade_generator.pth'  # path to checkpoint of G
+    checkpoint_W: str = './downloaded_repo/pretrained_weights/liveportrait/base_models/warping_module.pth'  # path to checkpoint of W
+    checkpoint_S: str = './downloaded_repo/pretrained_weights/liveportrait/base_models/retargeting_models/stitching_retargeting_module.pth'  # path to checkpoint to S and R_eyes, R_lip
+
+    # EXPORTED PARAMS
+    flag_use_half_precision: bool = True
+    flag_crop_driving_video: bool = False
+    device_id: int = 0
+    flag_lip_zero: bool = False
+    flag_eye_retargeting: bool = False
+    flag_lip_retargeting: bool = False
+    flag_stitching: bool = False
+    flag_relative_motion: bool = False
+    flag_pasteback: bool = False
+    flag_do_crop: bool = False
+    flag_do_rot: bool = False
+    flag_force_cpu: bool = False 
+    flag_do_torch_compile: bool = False  
+
+    # NOT EXPORTED PARAMS
+    lip_zero_threshold: float = 0.03 # threshold for flag_lip_zero
+    anchor_frame: int = 0 # TO IMPLEMENT
+
+    input_shape: Tuple[int, int] = (256, 256)  # input shape
+    output_format: Literal['mp4', 'gif'] = 'mp4'  # output video format
+    crf: int = 15  # crf for output video
+    output_fps: int = 25 # default output fps
+
+    mask_crop: ndarray = cv2.imread(make_abs_path('../utils/resources/mask_template.png'), cv2.IMREAD_COLOR)
+    size_gif: int = 256 # default gif size, TO IMPLEMENT
+    source_max_dim: int = 1280 # the max dim of height and width of source image
+    source_division: int = 2 # make sure the height and width of source image can be divided by this number