Initial deployment: ACMDM Motion Generation

.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
+checkpoints/**/*.tar filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,83 @@
 ---
 title: ACMDM Motion Generation
-emoji: 🌖
-colorFrom: yellow
-colorTo: indigo
+emoji: 🎭
+colorFrom: blue
+colorTo: purple
 sdk: gradio
-sdk_version: 5.49.1
+sdk_version: 4.0.0
 app_file: app.py
 pinned: false
+license: mit
+hardware: gpu-t4-small
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ACMDM Motion Generation
+
+Generate human motion animations from text descriptions using the ACMDM (Absolute Coordinates Make Motion Generation Easy) model.
+
+## 🎯 Features
+
+- **Text-to-Motion Generation**: Create realistic human motion from natural language descriptions
+- **Batch Processing**: Generate multiple motions at once
+- **Auto-Length Estimation**: AI automatically determines optimal motion length
+- **Flexible Parameters**: Adjust CFG scale, motion length, and more
+- **Real-time Preview**: See your generated motions instantly
+
+## 🚀 Usage
+
+1. **Enter a text description** of the motion you want (e.g., "A person is running on a treadmill.")
+2. **Adjust parameters** (optional):
+   - Motion length (40-196 frames)
+   - CFG scale (controls text alignment)
+   - Auto-length estimation
+3. **Click "Generate Motion"**
+4. **View and download** your generated motion video
+
+## 📝 Example Prompts
+
+- "A person is running on a treadmill."
+- "Someone is doing jumping jacks."
+- "A person walks forward and then turns around."
+- "A person is dancing energetically."
+
+## ⚙️ Parameters
+
+- **Motion Length**: Number of frames (40-196). Automatically rounded to multiples of 4.
+- **CFG Scale**: Classifier-free guidance scale (1.0-10.0). Higher = more text-aligned, Lower = more diverse.
+- **Auto-length**: Let AI estimate the optimal motion length based on your text.
+
+## 🔧 Technical Details
+
+This space uses pre-trained ACMDM models:
+- **Autoencoder**: AE_2D_Causal
+- **Diffusion Model**: ACMDM_Flow_S_PatchSize22
+- **Dataset**: HumanML3D (t2m)
+
+## 📚 Paper
+
+[Absolute Coordinates Make Motion Generation Easy](https://arxiv.org/abs/2505.19377)
+
+## 🤝 Citation
+
+```bibtex
+@article{meng2025absolute,
+    title={Absolute Coordinates Make Motion Generation Easy},
+    author={Meng, Zichong and Han, Zeyu and Peng, Xiaogang and Xie, Yiming and Jiang, Huaizu},
+    journal={arXiv preprint arXiv:2505.19377},
+    year={2025}
+}
+```
+
+## ⚠️ Notes
+
+- First generation may take 30-60 seconds (model loading)
+- Subsequent generations are faster (5-15 seconds)
+- GPU recommended for best performance
+- Works on CPU but slower
+
+## 🔗 Links
+
+- [GitHub Repository](https://github.com/neu-vi/ACMDM)
+- [Project Page](https://neu-vi.github.io/ACMDM/)
+- [Paper](https://arxiv.org/abs/2505.19377)
+

app.py

@@ -0,0 +1,604 @@
+"""
+Gradio Web Interface for ACMDM Motion Generation
+Milestone 3: User-friendly web interface for text-to-motion generation
+"""
+
+import os
+import sys
+from os.path import join as pjoin
+import torch
+import numpy as np
+import gradio as gr
+from typing import Optional, Tuple, List
+import tempfile
+import random
+
+# Import from sample.py
+from models.AE_2D_Causal import AE_models
+from models.ACMDM import ACMDM_models
+from models.LengthEstimator import LengthEstimator
+from utils.back_process import back_process
+from utils.motion_process import plot_3d_motion, t2m_kinematic_chain
+
+# Global variables for model caching
+_models_cache = {
+    'ae': None,
+    'acmdm': None,
+    'length_estimator': None,
+    'stats': None,
+    'device': None,
+    'loaded': False
+}
+
+
+def set_seed(seed):
+    """Set random seed for reproducibility"""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.backends.cudnn.benchmark = False
+
+
+def load_models_cached(
+    gpu_id: int = 0,
+    model_name: str = 'ACMDM_Flow_S_PatchSize22',
+    ae_name: str = 'AE_2D_Causal',
+    ae_model: str = 'AE_Model',
+    model_type: str = 'ACMDM-Flow-S-PatchSize22',
+    dataset_name: str = 't2m',
+    checkpoints_dir: str = './checkpoints',
+    use_length_estimator: bool = False
+) -> Tuple[dict, str]:
+    """
+    Load models with caching to avoid reloading on every request.
+    Returns (models_dict, status_message)
+    """
+    global _models_cache
+    
+    # Check if models are already loaded
+    if _models_cache['loaded']:
+        return _models_cache, "Models already loaded (using cache)"
+    
+    try:
+        # Determine device
+        if gpu_id >= 0 and torch.cuda.is_available():
+            device = torch.device(f"cuda:{gpu_id}")
+        else:
+            device = torch.device("cpu")
+        
+        _models_cache['device'] = device
+        status_messages = [f"Using device: {device}"]
+        
+        # Load AE
+        status_messages.append("Loading AE model...")
+        ae = AE_models[ae_model](input_width=3)
+        ae_ckpt_path = pjoin(checkpoints_dir, dataset_name, ae_name, 'model', 'latest.tar')
+        
+        if not os.path.exists(ae_ckpt_path):
+            return None, f"Error: AE checkpoint not found at {ae_ckpt_path}"
+        
+        ae_ckpt = torch.load(ae_ckpt_path, map_location='cpu')
+        ae.load_state_dict(ae_ckpt['ae'])
+        ae.eval()
+        ae.to(device)
+        _models_cache['ae'] = ae
+        status_messages.append(f"✓ Loaded AE from {ae_ckpt_path}")
+        
+        # Load ACMDM
+        status_messages.append("Loading ACMDM model...")
+        acmdm = ACMDM_models[model_type](input_dim=ae.output_emb_width, cond_mode='text')
+        acmdm_ckpt_path = pjoin(checkpoints_dir, dataset_name, model_name, 'model', 'latest.tar')
+        
+        if not os.path.exists(acmdm_ckpt_path):
+            return None, f"Error: ACMDM checkpoint not found at {acmdm_ckpt_path}"
+        
+        acmdm_ckpt = torch.load(acmdm_ckpt_path, map_location='cpu')
+        missing_keys, unexpected_keys = acmdm.load_state_dict(acmdm_ckpt['ema_acmdm'], strict=False)
+        assert len(unexpected_keys) == 0
+        assert all([k.startswith('clip_model.') for k in missing_keys])
+        acmdm.eval()
+        acmdm.to(device)
+        _models_cache['acmdm'] = acmdm
+        status_messages.append(f"✓ Loaded ACMDM from {acmdm_ckpt_path}")
+        
+        # Load LengthEstimator if needed
+        length_estimator = None
+        if use_length_estimator:
+            status_messages.append("Loading LengthEstimator...")
+            length_estimator = LengthEstimator(input_size=512, output_size=1)
+            length_estimator_path = pjoin(checkpoints_dir, dataset_name, 'length_estimator', 'model', 'latest.tar')
+            if os.path.exists(length_estimator_path):
+                length_ckpt = torch.load(length_estimator_path, map_location='cpu')
+                length_estimator.load_state_dict(length_ckpt['model'])
+                length_estimator.eval()
+                length_estimator.to(device)
+                _models_cache['length_estimator'] = length_estimator
+                status_messages.append(f"✓ Loaded LengthEstimator")
+            else:
+                status_messages.append(f"⚠ LengthEstimator not found, will use default length")
+        
+        # Load normalization stats
+        status_messages.append("Loading normalization statistics...")
+        after_mean = np.load(pjoin(checkpoints_dir, dataset_name, ae_name, 'AE_2D_Causal_Post_Mean.npy'))
+        after_std = np.load(pjoin(checkpoints_dir, dataset_name, ae_name, 'AE_2D_Causal_Post_Std.npy'))
+        joint_mean = np.load(f'utils/22x3_mean_std/{dataset_name}/22x3_mean.npy')
+        joint_std = np.load(f'utils/22x3_mean_std/{dataset_name}/22x3_std.npy')
+        eval_mean = np.load(f'utils/eval_mean_std/{dataset_name}/eval_mean.npy')
+        eval_std = np.load(f'utils/eval_mean_std/{dataset_name}/eval_std.npy')
+        
+        _models_cache['stats'] = {
+            'after_mean': after_mean,
+            'after_std': after_std,
+            'joint_mean': joint_mean,
+            'joint_std': joint_std,
+            'eval_mean': eval_mean,
+            'eval_std': eval_std
+        }
+        
+        _models_cache['loaded'] = True
+        status_message = "\n".join(status_messages)
+        return _models_cache, status_message
+        
+    except Exception as e:
+        error_msg = f"Error loading models: {str(e)}"
+        import traceback
+        error_msg += f"\n\nTraceback:\n{traceback.format_exc()}"
+        return None, error_msg
+
+
+def estimate_motion_length(text: str, models_cache: dict) -> int:
+    """Estimate motion length from text using LengthEstimator"""
+    if models_cache['length_estimator'] is None:
+        return None
+    
+    device = models_cache['device']
+    acmdm = models_cache['acmdm']
+    length_estimator = models_cache['length_estimator']
+    
+    with torch.no_grad():
+        text_emb = acmdm.encode_text([text])
+        pred_length = length_estimator(text_emb)
+        pred_length = int(pred_length.item() * 4)
+        pred_length = ((pred_length + 2) // 4) * 4
+        pred_length = max(40, min(196, pred_length))
+    return pred_length
+
+
+def generate_motion_single(
+    text: str,
+    motion_length: Optional[int],
+    cfg_scale: float,
+    use_auto_length: bool,
+    gpu_id: int,
+    seed: int
+) -> Tuple[Optional[str], str]:
+    """
+    Generate a single motion from text.
+    Returns (video_path, status_message)
+    """
+    global _models_cache
+    
+    try:
+        set_seed(seed)
+        
+        # Load models if not cached
+        if not _models_cache['loaded']:
+            models_cache, load_msg = load_models_cached(
+                gpu_id=gpu_id,
+                use_length_estimator=use_auto_length
+            )
+            if models_cache is None:
+                return None, f"Failed to load models:\n{load_msg}"
+        else:
+            models_cache = _models_cache
+        
+        device = models_cache['device']
+        ae = models_cache['ae']
+        acmdm = models_cache['acmdm']
+        stats = models_cache['stats']
+        
+        # Estimate length if needed
+        if motion_length is None or (motion_length == 0 and use_auto_length):
+            if use_auto_length and models_cache['length_estimator'] is not None:
+                motion_length = estimate_motion_length(text, models_cache)
+                status_msg = f"Estimated motion length: {motion_length} frames\n"
+            else:
+                motion_length = 120  # Default
+                status_msg = f"Using default motion length: {motion_length} frames\n"
+        else:
+            # Round to multiple of 4
+            motion_length = ((motion_length + 2) // 4) * 4
+            status_msg = f"Using specified motion length: {motion_length} frames\n"
+        
+        status_msg += f"Generating motion for: '{text}'...\n"
+        
+        # Generate motion
+        with torch.no_grad():
+            latent_length = motion_length // 4
+            m_lens = torch.tensor([latent_length], device=device)
+            pred_latents = acmdm.generate([text], m_lens, cfg_scale)
+            
+            # Denormalize latents
+            pred_latents_np = pred_latents.permute(0, 2, 3, 1).detach().cpu().numpy()
+            pred_latents_np = pred_latents_np * stats['after_std'] + stats['after_mean']
+            pred_latents_tensor = torch.from_numpy(pred_latents_np).to(device)
+            
+            # Decode through AE
+            pred_motions = ae.decode(pred_latents_tensor.permute(0, 3, 1, 2))
+            
+            # Denormalize motions
+            pred_motions_np = pred_motions.permute(0, 2, 3, 1).detach().cpu().numpy()
+            if stats['joint_mean'].ndim == 1:
+                pred_motions_np = pred_motions_np * stats['joint_std'][np.newaxis, np.newaxis, :, np.newaxis] + stats['joint_mean'][np.newaxis, np.newaxis, :, np.newaxis]
+            else:
+                pred_motions_np = pred_motions_np * stats['joint_std'][np.newaxis, ..., np.newaxis] + stats['joint_mean'][np.newaxis, ..., np.newaxis]
+            
+            # Back process to get RIC format, then recover joint positions
+            from utils.motion_process import recover_from_ric
+            
+            motion = pred_motions_np[0]  # (22, 3, seq_len)
+            motion = motion[:, :, :motion_length].transpose(2, 0, 1)  # (seq_len, 22, 3)
+            ric_data = back_process(motion, is_mesh=False)
+            ric_tensor = torch.from_numpy(ric_data).float()
+            joints = recover_from_ric(ric_tensor, joints_num=22).numpy()  # (seq_len, 22, 3)
+        
+        # Create temporary file for video
+        temp_dir = tempfile.mkdtemp()
+        video_path = pjoin(temp_dir, 'motion.mp4')
+        
+        # Generate video
+        plot_3d_motion(
+            video_path,
+            t2m_kinematic_chain,
+            joints,
+            title=text,
+            fps=20,
+            radius=4
+        )
+        
+        status_msg += "✓ Motion generated successfully!"
+        return video_path, status_msg
+        
+    except Exception as e:
+        error_msg = f"Error generating motion: {str(e)}"
+        import traceback
+        error_msg += f"\n\nTraceback:\n{traceback.format_exc()}"
+        return None, error_msg
+
+
+def generate_motion_batch(
+    text_file_content: str,
+    cfg_scale: float,
+    use_auto_length: bool,
+    gpu_id: int,
+    seed: int
+) -> Tuple[List[Optional[str]], str]:
+    """
+    Generate motions from a batch of text prompts.
+    Returns (list_of_video_paths, status_message)
+    """
+    # Parse text file
+    text_prompts = []
+    for line in text_file_content.strip().split('\n'):
+        line = line.strip()
+        if not line:
+            continue
+        if '#' in line:
+            parts = line.split('#')
+            text = parts[0].strip()
+            length_str = parts[1].strip() if len(parts) > 1 else 'NA'
+        else:
+            text = line
+            length_str = 'NA'
+        
+        if length_str.upper() == 'NA':
+            motion_length = None if use_auto_length else 120
+        else:
+            try:
+                motion_length = int(length_str)
+                motion_length = ((motion_length + 2) // 4) * 4
+            except:
+                motion_length = None if use_auto_length else 120
+        
+        text_prompts.append((text, motion_length))
+    
+    if not text_prompts:
+        return [], "No valid prompts found in file"
+    
+    status_msg = f"Processing {len(text_prompts)} prompts...\n"
+    video_paths = []
+    
+    for idx, (text, motion_length) in enumerate(text_prompts):
+        video_path, gen_msg = generate_motion_single(
+            text, motion_length, cfg_scale, use_auto_length, gpu_id, seed + idx
+        )
+        video_paths.append(video_path)
+        status_msg += f"\n[{idx+1}/{len(text_prompts)}] {gen_msg}\n"
+    
+    return video_paths, status_msg
+
+
+# Gradio Interface
+def create_interface():
+    """Create and configure the Gradio interface"""
+    
+    with gr.Blocks(title="ACMDM Motion Generation", theme=gr.themes.Soft()) as app:
+        gr.Markdown("""
+        # 🎭 ACMDM Motion Generation
+        
+        Generate human motion from text descriptions using the ACMDM (Absolute Coordinates Make Motion Generation Easy) model.
+        
+        **How to use:**
+        1. Enter a text description of the motion you want to generate
+        2. Adjust motion length (or use auto-estimate)
+        3. Click "Generate Motion" to create the animation
+        4. View and download the generated video
+        
+        **Example prompts:**
+        - "A person is running on a treadmill."
+        - "Someone is doing jumping jacks."
+        - "A person walks forward and then turns around."
+        """)
+        
+        with gr.Tabs():
+            # Single Generation Tab
+            with gr.Tab("Single Motion Generation"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        text_input = gr.Textbox(
+                            label="Motion Description",
+                            placeholder="A person is running on a treadmill.",
+                            lines=3,
+                            value="A person is running on a treadmill."
+                        )
+                        
+                        with gr.Row():
+                            motion_length = gr.Slider(
+                                label="Motion Length (frames)",
+                                minimum=40,
+                                maximum=196,
+                                value=120,
+                                step=4,
+                                info="Will be rounded to nearest multiple of 4"
+                            )
+                            use_auto_length = gr.Checkbox(
+                                label="Auto-estimate length",
+                                value=False,
+                                info="Use length estimator (ignores manual length if checked)"
+                            )
+                        
+                        cfg_scale = gr.Slider(
+                            label="CFG Scale",
+                            minimum=1.0,
+                            maximum=10.0,
+                            value=3.0,
+                            step=0.5,
+                            info="Classifier-free guidance scale (higher = more aligned with text)"
+                        )
+                        
+                        with gr.Row():
+                            gpu_id = gr.Number(
+                                label="GPU ID",
+                                value=0,
+                                precision=0,
+                                info="Use -1 for CPU"
+                            )
+                            seed = gr.Number(
+                                label="Random Seed",
+                                value=3407,
+                                precision=0
+                            )
+                        
+                        generate_btn = gr.Button("Generate Motion", variant="primary", size="lg")
+                    
+                    with gr.Column(scale=1):
+                        video_output = gr.Video(
+                            label="Generated Motion",
+                            format="mp4"
+                        )
+                        status_output = gr.Textbox(
+                            label="Status",
+                            lines=10,
+                            interactive=False
+                        )
+                
+                # Update model status after generation
+                def generate_and_update_status(text, motion_length, cfg_scale, use_auto_length, gpu_id, seed):
+                    video_path, status_msg = generate_motion_single(
+                        text, motion_length, cfg_scale, use_auto_length, gpu_id, seed
+                    )
+                    # Return video and status, plus trigger status update
+                    return video_path, status_msg
+                
+                generate_btn.click(
+                    fn=generate_and_update_status,
+                    inputs=[text_input, motion_length, cfg_scale, use_auto_length, gpu_id, seed],
+                    outputs=[video_output, status_output]
+                )
+            
+            # Batch Generation Tab
+            with gr.Tab("Batch Generation"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        batch_text_input = gr.Textbox(
+                            label="Text Prompts (one per line, format: text#length or text#NA)",
+                            placeholder="A person is running on a treadmill.#120\nSomeone is doing jumping jacks.#NA",
+                            lines=10,
+                            info="Each line: 'text#length' or 'text#NA' for auto-estimate"
+                        )
+                        
+                        batch_cfg_scale = gr.Slider(
+                            label="CFG Scale",
+                            minimum=1.0,
+                            maximum=10.0,
+                            value=3.0,
+                            step=0.5
+                        )
+                        
+                        batch_use_auto_length = gr.Checkbox(
+                            label="Auto-estimate length for NA",
+                            value=True
+                        )
+                        
+                        batch_gpu_id = gr.Number(
+                            label="GPU ID",
+                            value=0,
+                            precision=0
+                        )
+                        
+                        batch_seed = gr.Number(
+                            label="Random Seed",
+                            value=3407,
+                            precision=0
+                        )
+                        
+                        batch_generate_btn = gr.Button("Generate Batch", variant="primary", size="lg")
+                    
+                    with gr.Column(scale=1):
+                        batch_status_output = gr.Textbox(
+                            label="Batch Status",
+                            lines=15,
+                            interactive=False
+                        )
+                        batch_video_gallery = gr.Gallery(
+                            label="Generated Motions",
+                            show_label=True,
+                            elem_id="gallery",
+                            columns=2,
+                            rows=2,
+                            height="auto"
+                        )
+                
+                batch_generate_btn.click(
+                    fn=generate_motion_batch,
+                    inputs=[batch_text_input, batch_cfg_scale, batch_use_auto_length, batch_gpu_id, batch_seed],
+                    outputs=[batch_video_gallery, batch_status_output]
+                )
+            
+            # Model Management Tab
+            with gr.Tab("Model Management"):
+                with gr.Row():
+                    with gr.Column():
+                        model_status = gr.Textbox(
+                            label="Model Status",
+                            lines=15,
+                            interactive=False,
+                            value="⏳ Models not loaded yet. They will be loaded automatically on first generation."
+                        )
+                        with gr.Row():
+                            refresh_status_btn = gr.Button("🔄 Refresh Status", variant="primary")
+                            reload_models_btn = gr.Button("🗑️ Clear Cache", variant="secondary")
+                        
+                        gr.Markdown("""
+                        **Model Configuration:**
+                        - Model Name: `ACMDM_Flow_S_PatchSize22`
+                        - AE Name: `AE_2D_Causal`
+                        - Dataset: `t2m`
+                        - Checkpoints Directory: `./checkpoints`
+                        """)
+                
+                def check_model_status():
+                    """Check and display current model status"""
+                    global _models_cache
+                    if _models_cache['loaded']:
+                        device = _models_cache['device']
+                        status_lines = [
+                            "✓ MODELS LOADED AND READY",
+                            "=" * 50,
+                            f"📱 Device: {device}",
+                            f"💾 CUDA Available: {torch.cuda.is_available()}",
+                        ]
+                        
+                        if device.type == 'cuda':
+                            status_lines.append(f"🎮 GPU: {torch.cuda.get_device_name(device.index if device.index is not None else 0)}")
+                            status_lines.append(f"💾 GPU Memory: {torch.cuda.get_device_properties(device.index if device.index is not None else 0).total_memory / 1e9:.2f} GB")
+                        
+                        status_lines.extend([
+                            "",
+                            "📦 Loaded Models:",
+                            "  ✓ Autoencoder (AE_2D_Causal)",
+                            "  ✓ ACMDM Diffusion Model",
+                        ])
+                        
+                        if _models_cache['length_estimator'] is not None:
+                            status_lines.append("  ✓ Length Estimator")
+                        else:
+                            status_lines.append("  ⚠ Length Estimator (not loaded)")
+                        
+                        status_lines.extend([
+                            "",
+                            "📊 Statistics Loaded:",
+                            "  ✓ Post-AE Mean/Std",
+                            "  ✓ Joint Mean/Std",
+                            "  ✓ Eval Mean/Std",
+                            "",
+                            "✨ Models are cached and ready for generation!",
+                            "💡 Tip: Use 'Clear Cache' to force reload on next generation."
+                        ])
+                        
+                        return "\n".join(status_lines)
+                    else:
+                        return (
+                            "⏳ Models not loaded yet. They will be loaded automatically on first generation.\n"
+                            "📝 To load models now, go to 'Single Motion Generation' tab and click 'Generate Motion'.\n"
+                            "⏱️ First generation will take 30-60 seconds (model loading time).\n"
+                            "⚡ Subsequent generations will be much faster (5-15 seconds)."
+                        )
+                
+                def reload_models():
+                    """Clear model cache"""
+                    global _models_cache
+                    _models_cache['loaded'] = False
+                    _models_cache['ae'] = None
+                    _models_cache['acmdm'] = None
+                    _models_cache['length_estimator'] = None
+                    _models_cache['stats'] = None
+                    _models_cache['device'] = None
+                    return (
+                        "🗑️ Model cache cleared. Models will be automatically reloaded on your next generation request.\n"
+                        "💡 Click 'Refresh Status' after generating to see updated status."
+                    )
+                
+                # Button callbacks
+                refresh_status_btn.click(
+                    fn=check_model_status,
+                    outputs=model_status
+                )
+                
+                reload_models_btn.click(
+                    fn=reload_models,
+                    outputs=model_status
+                )
+                
+                # Update status on tab load
+                app.load(
+                    fn=check_model_status,
+                    outputs=model_status
+                )
+        
+        gr.Markdown("""
+        ---
+        **Note:** First generation may take longer as models need to be loaded. Subsequent generations will be faster.
+        
+        **Tips:**
+        - Use descriptive text prompts for better results
+        - Adjust CFG scale: higher values (3-5) for more text alignment, lower values (1-2) for more diversity
+        - Motion length should be a multiple of 4 (automatically rounded)
+        - For batch processing, use the format: `text description#length` or `text description#NA`
+        """)
+    
+    return app
+
+
+if __name__ == "__main__":
+    # Create and launch the interface
+    app = create_interface()
+    
+    # Launch with sharing option
+    app.launch(
+        server_name="0.0.0.0",  # Allow external connections
+        server_port=7860,        # Default Gradio port
+        share=False,             # Set to True for public link (requires ngrok)
+        show_error=True
+    )
+

checkpoints/t2m/ACMDM_Flow_S_PatchSize22/model/latest.tar

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

checkpoints/t2m/AE_2D_Causal/AE_2D_Causal_Post_Mean.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:49354ec5e7cfc75121d3d4ddf323364debd86e8815e9a153176c9d538b1fbf17
+size 144

checkpoints/t2m/AE_2D_Causal/AE_2D_Causal_Post_Std.npy

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

checkpoints/t2m/AE_2D_Causal/model/latest.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c9601098c69561fcc12a7e91762c3caa03e5ab03e2c0488f44b2c4edbd08ee3
+size 204997782

checkpoints/t2m/length_estimator/model/finest.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90512e0e893186e3b09b24b006bfbf51f1ad71ac4c6626c9b1309373db675d12
+size 1745581

requirements.txt

@@ -0,0 +1,30 @@
+# Requirements for HuggingFace Spaces Deployment
+# This file is used by HuggingFace Spaces to install dependencies
+
+# Core ML libraries
+torch>=2.0.0
+torchvision>=0.15.0
+numpy>=1.21.0
+scipy>=1.7.0
+
+# Gradio for web interface
+gradio>=4.0.0
+
+# CLIP for text encoding
+git+https://github.com/openai/CLIP.git
+
+# Additional dependencies
+einops>=0.6.0
+timm>=0.6.0
+tqdm>=4.64.0
+matplotlib>=3.5.0
+opencv-python>=4.6.0
+Pillow>=9.0.0
+
+# For model loading and processing
+h5py>=3.7.0
+scikit-learn>=1.0.0
+
+# Utilities
+gdown>=4.6.0
+

utils/22x3_mean_std/t2m/22x3_mean.npy

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

utils/22x3_mean_std/t2m/22x3_std.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4733059903169bc95993999fde322d8e84fb95f3e6ff1b277c283a0eb877d0c5
+size 140

utils/back_process.py

@@ -0,0 +1,255 @@
+from utils.skeleton import Skeleton
+from utils.quaternion import *
+from utils.motion_process import t2m_kinematic_chain, t2m_raw_offsets
+
+import torch
+import os
+
+n_raw_offsets = torch.from_numpy(t2m_raw_offsets)
+kinematic_chain = t2m_kinematic_chain
+l_idx1, l_idx2 = 5, 8
+face_joint_indx = [2, 1, 17, 16]
+
+# Lazy loading of tgt_offsets - only needed when is_mesh=True
+_tgt_offsets_cache = None
+
+def _get_tgt_offsets():
+    """Lazily load target offsets for mesh processing. Only called when is_mesh=True."""
+    global _tgt_offsets_cache
+    if _tgt_offsets_cache is None:
+        example_data_path = os.path.join("datasets/HumanML3D/new_joints", "000021" + '.npy')
+        if not os.path.exists(example_data_path):
+            raise FileNotFoundError(
+                f"Example data file not found: {example_data_path}\n"
+                "This file is only needed for mesh-level motion generation (is_mesh=True).\n"
+                "For regular text-to-motion generation, use is_mesh=False."
+            )
+        example_data = np.load(example_data_path)
+        example_data = example_data.reshape(len(example_data), -1, 3)
+        example_data = torch.from_numpy(example_data)
+        tgt_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')
+        _tgt_offsets_cache = tgt_skel.get_offsets_joints(example_data[0])
+    return _tgt_offsets_cache
+
+l_idx1, l_idx2 = 5, 8
+fid_r, fid_l = [8, 11], [7, 10]
+r_hip, l_hip = 2, 1
+joints_num = 22
+
+def uniform_skeleton(positions, target_offset):
+    src_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')
+    src_offset = src_skel.get_offsets_joints(torch.from_numpy(positions[0]))
+    src_offset = src_offset.numpy()
+    tgt_offset = target_offset.numpy()
+    # print(src_offset)
+    # print(tgt_offset)
+    '''Calculate Scale Ratio as the ratio of legs'''
+    src_leg_len = np.abs(src_offset[l_idx1]).max() + np.abs(src_offset[l_idx2]).max()
+    tgt_leg_len = np.abs(tgt_offset[l_idx1]).max() + np.abs(tgt_offset[l_idx2]).max()
+
+    scale_rt = tgt_leg_len / src_leg_len
+    # print(scale_rt)
+    src_root_pos = positions[:, 0]
+    tgt_root_pos = src_root_pos * scale_rt
+
+    '''Inverse Kinematics'''
+    quat_params = src_skel.inverse_kinematics_np(positions, face_joint_indx)
+    # print(quat_params.shape)
+
+    '''Forward Kinematics'''
+    src_skel.set_offset(target_offset)
+    new_joints = src_skel.forward_kinematics_np(quat_params, tgt_root_pos)
+    return new_joints
+
+
+def process_file(positions, feet_thre, is_mesh=False):
+    # (seq_len, joints_num, 3)
+    #     '''Down Sample'''
+    #     positions = positions[::ds_num]
+
+    if is_mesh:
+        '''Uniform Skeleton'''
+        tgt_offsets = _get_tgt_offsets()  # Lazy load only when needed
+        positions = uniform_skeleton(positions, tgt_offsets)
+
+        '''Put on Floor'''
+        floor_height = positions.min(axis=0).min(axis=0)[1]
+        positions[:, :, 1] -= floor_height
+        #     print(floor_height)
+
+        #     plot_3d_motion("./positions_1.mp4", kinematic_chain, positions, 'title', fps=20)
+
+        '''XZ at origin'''
+        root_pos_init = positions[0]
+        root_pose_init_xz = root_pos_init[0] * np.array([1, 0, 1])
+        positions = positions - root_pose_init_xz
+
+        # '''Move the first pose to origin '''
+        # root_pos_init = positions[0]
+        # positions = positions - root_pos_init[0]
+
+        '''All initially face Z+'''
+        r_hip, l_hip, sdr_r, sdr_l = face_joint_indx
+        across1 = root_pos_init[r_hip] - root_pos_init[l_hip]
+        across2 = root_pos_init[sdr_r] - root_pos_init[sdr_l]
+        across = across1 + across2
+        across = across / np.sqrt((across ** 2).sum(axis=-1))[..., np.newaxis]
+
+        # forward (3,), rotate around y-axis
+        forward_init = np.cross(np.array([[0, 1, 0]]), across, axis=-1)
+        # forward (3,)
+        forward_init = forward_init / np.sqrt((forward_init ** 2).sum(axis=-1))[..., np.newaxis]
+
+        #     print(forward_init)
+
+        target = np.array([[0, 0, 1]])
+        root_quat_init = qbetween_np(forward_init, target)
+        root_quat_init = np.ones(positions.shape[:-1] + (4,)) * root_quat_init
+
+        positions_b = positions.copy()
+
+        positions = qrot_np(root_quat_init, positions)
+
+    #     plot_3d_motion("./positions_2.mp4", kinematic_chain, positions, 'title', fps=20)
+
+    '''New ground truth positions'''
+    global_positions = positions.copy()
+
+    # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*')
+    # plt.plot(positions[:, 0, 0], positions[:, 0, 2], marker='o', color='r')
+    # plt.xlabel('x')
+    # plt.ylabel('z')
+    # plt.axis('equal')
+    # plt.show()
+
+    """ Get Foot Contacts """
+
+    def foot_detect(positions, thres):
+        velfactor, heightfactor = np.array([thres, thres]), np.array([3.0, 2.0])
+
+        feet_l_x = (positions[1:, fid_l, 0] - positions[:-1, fid_l, 0]) ** 2
+        feet_l_y = (positions[1:, fid_l, 1] - positions[:-1, fid_l, 1]) ** 2
+        feet_l_z = (positions[1:, fid_l, 2] - positions[:-1, fid_l, 2]) ** 2
+        #     feet_l_h = positions[:-1,fid_l,1]
+        #     feet_l = (((feet_l_x + feet_l_y + feet_l_z) < velfactor) & (feet_l_h < heightfactor)).astype(np.float)
+        feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).astype(np.float32)
+
+        feet_r_x = (positions[1:, fid_r, 0] - positions[:-1, fid_r, 0]) ** 2
+        feet_r_y = (positions[1:, fid_r, 1] - positions[:-1, fid_r, 1]) ** 2
+        feet_r_z = (positions[1:, fid_r, 2] - positions[:-1, fid_r, 2]) ** 2
+        #     feet_r_h = positions[:-1,fid_r,1]
+        #     feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor) & (feet_r_h < heightfactor)).astype(np.float)
+        feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor)).astype(np.float32)
+        return feet_l, feet_r
+
+    #
+    feet_l, feet_r = foot_detect(positions, feet_thre)
+    # feet_l, feet_r = foot_detect(positions, 0.002)
+
+    '''Quaternion and Cartesian representation'''
+    r_rot = None
+
+    def get_rifke(positions):
+        '''Local pose'''
+        positions[..., 0] -= positions[:, 0:1, 0]
+        positions[..., 2] -= positions[:, 0:1, 2]
+        '''All pose face Z+'''
+        positions = qrot_np(np.repeat(r_rot[:, None], positions.shape[1], axis=1), positions)
+        return positions
+
+    def get_quaternion(positions):
+        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
+        # (seq_len, joints_num, 4)
+        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=False)
+
+        '''Fix Quaternion Discontinuity'''
+        quat_params = qfix(quat_params)
+        # (seq_len, 4)
+        r_rot = quat_params[:, 0].copy()
+        #     print(r_rot[0])
+        '''Root Linear Velocity'''
+        # (seq_len - 1, 3)
+        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
+        #     print(r_rot.shape, velocity.shape)
+        velocity = qrot_np(r_rot[1:], velocity)
+        '''Root Angular Velocity'''
+        # (seq_len - 1, 4)
+        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
+        quat_params[1:, 0] = r_velocity
+        # (seq_len, joints_num, 4)
+        return quat_params, r_velocity, velocity, r_rot
+
+    def get_cont6d_params(positions):
+        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
+        # (seq_len, joints_num, 4)
+        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=True)
+
+        '''Quaternion to continuous 6D'''
+        cont_6d_params = quaternion_to_cont6d_np(quat_params)
+        # (seq_len, 4)
+        r_rot = quat_params[:, 0].copy()
+        #     print(r_rot[0])
+        '''Root Linear Velocity'''
+        # (seq_len - 1, 3)
+        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
+        #     print(r_rot.shape, velocity.shape)
+        velocity = qrot_np(r_rot[1:], velocity)
+        '''Root Angular Velocity'''
+        # (seq_len - 1, 4)
+        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
+        # (seq_len, joints_num, 4)
+        return cont_6d_params, r_velocity, velocity, r_rot
+
+    cont_6d_params, r_velocity, velocity, r_rot = get_cont6d_params(positions)
+    positions = get_rifke(positions)
+
+    #     trejec = np.cumsum(np.concatenate([np.array([[0, 0, 0]]), velocity], axis=0), axis=0)
+    #     r_rotations, r_pos = recover_ric_glo_np(r_velocity, velocity[:, [0, 2]])
+
+    # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*')
+    # plt.plot(ground_positions[:, 0, 0], ground_positions[:, 0, 2], marker='o', color='r')
+    # plt.plot(trejec[:, 0], trejec[:, 2], marker='^', color='g')
+    # plt.plot(r_pos[:, 0], r_pos[:, 2], marker='s', color='y')
+    # plt.xlabel('x')
+    # plt.ylabel('z')
+    # plt.axis('equal')
+    # plt.show()
+
+    '''Root height'''
+    root_y = positions[:, 0, 1:2]
+
+    '''Root rotation and linear velocity'''
+    # (seq_len-1, 1) rotation velocity along y-axis
+    # (seq_len-1, 2) linear velovity on xz plane
+    r_velocity = np.arcsin(r_velocity[:, 2:3])
+    l_velocity = velocity[:, [0, 2]]
+    #     print(r_velocity.shape, l_velocity.shape, root_y.shape)
+    root_data = np.concatenate([r_velocity, l_velocity, root_y[:-1]], axis=-1)
+
+    '''Get Joint Rotation Representation'''
+    # (seq_len, (joints_num-1) *6) quaternion for skeleton joints
+    rot_data = cont_6d_params[:, 1:].reshape(len(cont_6d_params), -1)
+
+    '''Get Joint Rotation Invariant Position Represention'''
+    # (seq_len, (joints_num-1)*3) local joint position
+    ric_data = positions[:, 1:].reshape(len(positions), -1)
+
+    '''Get Joint Velocity Representation'''
+    # (seq_len-1, joints_num*3)
+    local_vel = qrot_np(np.repeat(r_rot[:-1, None], global_positions.shape[1], axis=1),
+                        global_positions[1:] - global_positions[:-1])
+    local_vel = local_vel.reshape(len(local_vel), -1)
+
+    data = root_data
+    data = np.concatenate([data, ric_data[:-1]], axis=-1)
+    data = np.concatenate([data, rot_data[:-1]], axis=-1)
+    #     print(data.shape, local_vel.shape)
+    data = np.concatenate([data, local_vel], axis=-1)
+    data = np.concatenate([data, feet_l, feet_r], axis=-1)
+
+    return data, global_positions, positions, l_velocity
+
+
+def back_process(data, is_mesh=False):
+    data, ground_positions, positions, l_velocity = process_file(data, 0.002, is_mesh=is_mesh)
+    return data[:, :67]

utils/cal_ae_post_mean_std.py

@@ -0,0 +1,64 @@
+import numpy as np
+import os
+from os.path import join as pjoin
+from tqdm import tqdm
+import torch
+import argparse
+from models.AE_2D_Causal import AE_models
+
+#################################################################################
+#                           Calculate Post AE/VAE Mean Std                      #
+#################################################################################
+
+def mean_variance(data_dir, save_dir, ae, tp='AE'):
+    file_list = os.listdir(data_dir)
+    data_list = []
+    mean = np.load(f'utils/22x3_mean_std/t2m/22x3_mean.npy')
+    std = np.load(f'utils/22x3_mean_std/t2m/22x3_std.npy')
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    ae = ae.to(device)
+
+    for file in tqdm(file_list):
+        data = np.load(pjoin(data_dir, file))
+        if len(data.shape) == 2:
+            data = np.expand_dims(data, axis=0)
+        if np.isnan(data).any():
+            print(file)
+            continue
+        data = data[:(data.shape[0]//4)*4, :, :]
+        if data.shape[0] == 0:
+            continue
+        data = (data - mean) / std
+        data = torch.from_numpy(data).to(device)
+        with torch.no_grad():
+            data_list.append(ae.encode(data.unsqueeze(0)).squeeze(0).cpu().numpy())
+
+    data = np.concatenate(data_list, axis=1)
+    data = data.reshape(4, -1)
+    print('Data Range:')
+    print(data.min(),data.max())
+    Mean = data.mean(axis=1)
+    Std = data.std(axis=1)
+    print('Mean/Std:')
+    print(Mean, Std)
+
+    np.save(pjoin(save_dir, f'{tp}_2D_Causal_Post_Mean.npy'), Mean)
+    np.save(pjoin(save_dir, f'{tp}_2D_Causal_Post_Std.npy'), Std)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    data_dir = 'datasets/HumanML3D/new_joints/'
+
+    parser.add_argument('--is_ae', action="store_true")
+    parser.add_argument('--ae_name', type=str, default="AE_2D_Causal")
+    args = parser.parse_args()
+
+    if args.is_ae:
+        ae = AE_models["AE_Model"](input_width=3)
+    else:
+        ae = AE_models["VAE_Model"](input_width=3)
+    ckpt = torch.load(f'checkpoints/t2m/{args.ae_name}/model/latest.tar', map_location='cpu')
+    ae.load_state_dict(ckpt['ae'])
+    ae = ae.eval()
+    save_dir = f'checkpoints/t2m/{args.ae_name}'
+    mean_variance(data_dir, save_dir, ae, tp='AE' if args.is_ae else 'VAE')

utils/cal_mean_std.py

@@ -0,0 +1,35 @@
+import numpy as np
+import os
+from os.path import join as pjoin
+from tqdm import tqdm
+
+#################################################################################
+#                      Calculate Absolute Coordinate Mean Std                   #
+#################################################################################
+def mean_variance(data_dir, save_dir):
+    file_list = os.listdir(data_dir)
+    data_list = []
+
+    for file in tqdm(file_list):
+        data = np.load(pjoin(data_dir, file))
+        if len(data.shape) == 2:
+            data = np.expand_dims(data, axis=0)
+        if np.isnan(data).any():
+            print(file)
+            continue
+        data_list.append(data.reshape(-1, 3))
+
+    data = np.concatenate(data_list, axis=0)
+    print(data.shape)
+    Mean = data.mean(axis=0)
+    Std = data.std(axis=0)
+
+    np.save(pjoin(save_dir, 'Mean_22x3.npy'), Mean)
+    np.save(pjoin(save_dir, 'Std_22x3.npy'), Std)
+
+    return Mean, Std
+
+if __name__ == '__main__':
+    data_dir1 = 'datasets/HumanML3D/new_joints/'
+    save_dir1 = 'datasets/HumanML3D/'
+    mean, std = mean_variance(data_dir1, save_dir1)

utils/cal_mesh_ae_post_mean_std.py

@@ -0,0 +1,107 @@
+# very hard coded will clean and refine later
+import numpy as np
+import os
+from os.path import join as pjoin
+from tqdm import tqdm
+import torch
+from models.AE_Mesh import AE_models
+from human_body_prior.body_model.body_model import BodyModel
+
+def downsample(data_dir, save_dir):
+    ae = AE_models["AE_Model"](test_mode=True)
+    ckpt = torch.load('checkpoints/t2m/AE_Mesh/model/latest.tar', map_location='cpu')
+    model_key = 'ae'
+    ae.load_state_dict(ckpt[model_key])
+    ae = ae.eval()
+    ae = ae.cuda()
+
+    mean = np.load(f'utils/mesh_mean_std/t2m/mesh_mean.npy')  # or yours
+    std = np.load(f'utils/mesh_mean_std/t2m/mesh_std.npy')  # or yours
+
+    bm_path = './body_models/smplh/neutral/model.npz'
+    dmpl_path = './body_models/dmpls/neutral/model.npz'
+    num_betas = 10
+    num_dmpls = 8
+    bm = BodyModel(bm_fname=bm_path, num_betas=num_betas, num_dmpls=num_dmpls, dmpl_fname=dmpl_path).cuda()
+    comp_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    file_list = os.listdir(data_dir)
+    file_list.sort()
+    for file in tqdm(file_list):
+        data = np.load(pjoin(data_dir, file))
+
+        if np.isnan(data).any():
+            print(f"Skipping {file} due to NaN")
+            continue
+
+        body_parms = {
+            'root_orient': torch.from_numpy(data[:, :3]).to(comp_device).float(),
+            'pose_body': torch.from_numpy(data[:, 3:66]).to(comp_device).float(),
+            'pose_hand': torch.from_numpy(data[:, 66:52*3]).to(comp_device).float(),
+            'trans': torch.from_numpy(data[:, 52*3:53*3]).to(comp_device).float(),
+            'betas': torch.from_numpy(data[:, 53*3:53*3+10]).to(comp_device).float(),
+            'dmpls': torch.from_numpy(data[:, 53*3+10:]).to(comp_device).float()
+        }
+
+        with torch.no_grad():
+            verts = bm(**body_parms).v
+        verts[:, :, 1] -= verts[:, :, 1].min()
+        verts = verts.detach().cpu().numpy()
+        "Z Normalization"
+        vertss = (verts - mean) / std
+        T = vertss.shape[0]
+        data = torch.from_numpy(vertss).float().to(comp_device)
+        if T % 16 != 0:
+            pad_len = 16 - (T % 16)
+            pad_data = torch.zeros((pad_len, 6890, 3), dtype=data.dtype, device=comp_device)
+            data = torch.cat([data, pad_data], dim=0)
+
+        outputs = []
+        with torch.no_grad():
+            for i in range(0, data.shape[0], 16):
+                chunk = data[i:i+16].unsqueeze(0)
+                out = ae.encode(chunk).squeeze(0).cpu().numpy()
+                outputs.append(out)
+        downsampled = np.concatenate(outputs, axis=0)[:T]
+
+        np.save(pjoin(save_dir, f"{file}"), downsampled)
+
+#################################################################################
+#                                Calculate Mean Std                             #
+#################################################################################
+def mean_variance(data_dir, save_dir):
+    file_list = os.listdir(data_dir)
+    data_list = []
+
+    for file in tqdm(file_list):
+        data = np.load(pjoin(data_dir, file))
+        if np.isnan(data).any():
+            print(file)
+            continue
+        data_list.append(data.reshape(-1, 12))
+
+    data = np.concatenate(data_list, axis=0)
+    print(data.shape)
+    Mean = data.mean(axis=0)
+    Std = data.std(axis=0)
+
+    np.save(pjoin(save_dir, 'AE_Mesh_Post_Mean.npy'), Mean)
+    np.save(pjoin(save_dir, 'AE_Mesh_Post_Std.npy'), Std)
+    print(data.min(),data.max())
+    Mean2 = data.mean()
+    Std2 = data.std()
+    print(Mean, Std)
+    print(Mean2, Std2)
+
+    return Mean, Std
+
+if __name__ == '__main__':
+    data_dir = 'datasets/HumanML3D/meshes/'
+    save_dir = 'datasets/HumanML3D/meshes_after_ae/'
+
+    os.makedirs(save_dir, exist_ok=True)
+    downsample(data_dir, save_dir)
+
+    data_dir1 = 'datasets/HumanML3D/meshes_after_ae/'
+    save_dir1 = 'checkpoints/t2m/AE_Mesh/'
+    mean, std = mean_variance(data_dir1, save_dir1)

utils/datasets.py

@@ -0,0 +1,452 @@
+from os.path import join as pjoin
+from torch.utils import data
+import numpy as np
+import torch
+from tqdm import tqdm
+from torch.utils.data._utils.collate import default_collate
+import random
+import codecs as cs
+from utils.glove import GloVe
+from human_body_prior.body_model.body_model import BodyModel
+
+#################################################################################
+#                                  Collate Function                             #
+#################################################################################
+def collate_fn(batch):
+    batch.sort(key=lambda x: x[3], reverse=True)
+    return default_collate(batch)
+
+#################################################################################
+#                                      Datasets                                 #
+#################################################################################
+class AEDataset(data.Dataset):
+    def __init__(self, mean, std, motion_dir, window_size, split_file):
+        self.data = []
+        self.lengths = []
+        id_list = []
+        with open(split_file, 'r') as f:
+            for line in f.readlines():
+                id_list.append(line.strip())
+
+        for name in tqdm(id_list):
+            try:
+                motion = np.load(pjoin(motion_dir, name + '.npy'))
+                if len(motion.shape) == 2: # B,L,J,3
+                    motion = np.expand_dims(motion, axis=0)
+                if motion.shape[0] < window_size:
+                    continue
+                self.lengths.append(motion.shape[0] - window_size)
+                self.data.append(motion)
+            except Exception as e:
+                pass
+        self.cumsum = np.cumsum([0] + self.lengths)
+        self.window_size = window_size
+
+        self.mean = mean
+        self.std = std
+        print("Total number of motions {}, snippets {}".format(len(self.data), self.cumsum[-1]))
+
+    def __len__(self):
+        return self.cumsum[-1]
+
+    def __getitem__(self, item):
+        if item != 0:
+            motion_id = np.searchsorted(self.cumsum, item) - 1
+            idx = item - self.cumsum[motion_id] - 1
+        else:
+            motion_id = 0
+            idx = 0
+        motion = self.data[motion_id][idx:idx + self.window_size]
+        "Z Normalization"
+        motion = (motion - self.mean) / self.std
+
+        return motion
+
+
+class AEMeshDataset(data.Dataset):
+    def __init__(self, mean, std, motion_dir, window_size, split_file):
+        self.data = []
+        self.lengths = []
+        id_list = []
+        with open(split_file, 'r') as f:
+            for line in f.readlines():
+                id_list.append(line.strip())
+
+        for name in id_list:
+            try:
+                motion = np.load(pjoin(motion_dir, name + '.npy'))
+                if motion.shape[0] < window_size:
+                    continue
+                self.lengths.append(motion.shape[0] - motion.shape[0]+1)
+                self.data.append(motion)
+            except Exception as e:
+                pass
+        self.cumsum = np.cumsum([0] + self.lengths)
+        self.window_size = window_size
+
+        self.mean = mean
+        self.std = std
+        num_betas = 10
+        num_dmpls = 8
+        self.bm = BodyModel(bm_fname='./body_models/smplh/neutral/model.npz', num_betas=num_betas, num_dmpls=num_dmpls, dmpl_fname='./body_models/dmpls/neutral/model.npz')
+        print("Total number of motions {}, snippets {}".format(len(self.data), self.cumsum[-1]))
+
+    def __len__(self):
+        return self.cumsum[-1]
+
+    def __getitem__(self, item):
+        motion = self.data[item]
+        body_parms = {
+            'root_orient': torch.from_numpy(motion[:, :3]).float(),
+            'pose_body': torch.from_numpy(motion[:, 3:66]).float(),
+            'pose_hand': torch.from_numpy(motion[:, 66:52*3]).float(),
+            'trans': torch.from_numpy(motion[:, 52*3:53*3]).float(),
+            'betas': torch.from_numpy(motion[:, 53*3:53*3+10]).float(),
+            'dmpls': torch.from_numpy(motion[:, 53*3+10:]).float()
+        }
+        body_parms['betas']= torch.zeros_like(torch.from_numpy(motion[:, 53*3:53*3+10]).float())
+        with torch.no_grad():
+            verts = self.bm(**body_parms).v
+        verts[:, :, 1] -= verts[:, :, 1].min()
+        idx = random.randint(0, verts.shape[0] - 1)
+        verts = verts[idx:idx + self.window_size].detach().cpu().numpy()
+        "Z Normalization"
+        verts = (verts - self.mean) / self.std
+
+        return verts
+
+
+
+class Text2MotionDataset(data.Dataset):
+    def __init__(self, mean, std, split_file, dataset_name, motion_dir, text_dir, unit_length, max_motion_length,
+                 max_text_length, evaluation=False, is_mesh=False):
+        self.evaluation = evaluation
+        self.max_length = 20
+        self.pointer = 0
+        self.max_motion_length = max_motion_length
+        self.max_text_len = max_text_length
+        self.unit_length = unit_length
+        min_motion_len = 40 if dataset_name =='t2m' else 24
+
+        data_dict = {}
+        id_list = []
+        with cs.open(split_file, 'r') as f:
+            for line in f.readlines():
+                id_list.append(line.strip())
+
+        new_name_list = []
+        length_list = []
+        for name in tqdm(id_list):
+            try:
+                motion = np.load(pjoin(motion_dir, name + '.npy'))
+                if len(motion.shape) == 2:
+                    motion = np.expand_dims(motion, axis=0)
+                if is_mesh:
+                    if (len(motion)) < min_motion_len:
+                        continue
+                else:
+                    if (len(motion)) < min_motion_len or (len(motion) >= 200):
+                        continue
+                text_data = []
+                flag = False
+                with cs.open(pjoin(text_dir, name + '.txt')) as f:
+                    for line in f.readlines():
+                        text_dict = {}
+                        line_split = line.strip().split('#')
+                        caption = line_split[0]
+                        tokens = line_split[1].split(' ')
+                        f_tag = float(line_split[2])
+                        to_tag = float(line_split[3])
+                        f_tag = 0.0 if np.isnan(f_tag) else f_tag
+                        to_tag = 0.0 if np.isnan(to_tag) else to_tag
+
+                        text_dict['caption'] = caption
+                        text_dict['tokens'] = tokens
+                        if f_tag == 0.0 and to_tag == 0.0:
+                            flag = True
+                            text_data.append(text_dict)
+                        else:
+                            try:
+                                n_motion = motion[int(f_tag*20) : int(to_tag*20)]
+                                if (len(n_motion)) < min_motion_len or (len(n_motion) >= 200):
+                                    continue
+                                new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                                while new_name in data_dict:
+                                    new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                                data_dict[new_name] = {'motion': n_motion,
+                                                       'length': len(n_motion),
+                                                       'text':[text_dict]}
+                                new_name_list.append(new_name)
+                                length_list.append(len(n_motion))
+                            except:
+                                print(line_split)
+                                print(line_split[2], line_split[3], f_tag, to_tag, name)
+
+                if flag:
+                    data_dict[name] = {'motion': motion,
+                                       'length': len(motion),
+                                       'text': text_data}
+                    new_name_list.append(name)
+                    length_list.append(len(motion))
+            except:
+                pass
+        if self.evaluation:
+            self.w_vectorizer = GloVe('./glove', 'our_vab')
+            name_list, length_list = zip(*sorted(zip(new_name_list, length_list), key=lambda x: x[1]))
+        else:
+            name_list, length_list = new_name_list, length_list
+        self.mean = mean
+        self.std = std
+        self.length_arr = np.array(length_list)
+        self.data_dict = data_dict
+        self.name_list = name_list
+        if self.evaluation:
+            self.reset_max_len(self.max_length)
+
+    def reset_max_len(self, length):
+        assert length <= self.max_motion_length
+        self.pointer = np.searchsorted(self.length_arr, length)
+        print("Pointer Pointing at %d"%self.pointer)
+        self.max_length = length
+
+    def transform(self, data, mean=None, std=None):
+        if mean is None and std is None:
+            return (data - self.mean) / self.std
+        else:
+            return (data - mean) / std
+
+    def inv_transform(self, data, mean=None, std=None):
+        if mean is None and std is None:
+            return data * self.std + self.mean
+        else:
+            return data * std + mean
+
+    def __len__(self):
+        return len(self.data_dict) - self.pointer
+
+    def __getitem__(self, item):
+        idx = self.pointer + item
+        data = self.data_dict[self.name_list[idx]]
+        motion, m_length, text_list = data['motion'], data['length'], data['text']
+        # Randomly select a caption
+        text_data = random.choice(text_list)
+        caption, tokens = text_data['caption'], text_data['tokens']
+
+        if self.evaluation:
+            if len(tokens) < self.max_text_len:
+                # pad with "unk"
+                tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+                sent_len = len(tokens)
+                tokens = tokens + ['unk/OTHER'] * (self.max_text_len + 2 - sent_len)
+            else:
+                # crop
+                tokens = tokens[:self.max_text_len]
+                tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+                sent_len = len(tokens)
+            pos_one_hots = []
+            word_embeddings = []
+            for token in tokens:
+                word_emb, pos_oh = self.w_vectorizer[token]
+                pos_one_hots.append(pos_oh[None, :])
+                word_embeddings.append(word_emb[None, :])
+            pos_one_hots = np.concatenate(pos_one_hots, axis=0)
+            word_embeddings = np.concatenate(word_embeddings, axis=0)
+
+        if self.unit_length < 10:
+            coin2 = np.random.choice(['single', 'single', 'double'])
+        else:
+            coin2 = 'single'
+
+        if coin2 == 'double':
+            m_length = (m_length // self.unit_length - 1) * self.unit_length
+        elif coin2 == 'single':
+            m_length = (m_length // self.unit_length) * self.unit_length
+        idx = random.randint(0, len(motion) - m_length)
+        motion = motion[idx:idx+m_length]
+
+        "Z Normalization"
+        motion = (motion - self.mean) / self.std
+
+        if m_length < self.max_motion_length:
+            motion = np.concatenate([motion,
+                                         np.zeros((self.max_motion_length - m_length, motion.shape[1], motion.shape[2]))
+                                         ], axis=0)
+        elif m_length > self.max_motion_length:
+            idx = random.randint(0, m_length - self.max_motion_length)
+            motion = motion[idx:idx + self.max_motion_length]
+        if self.evaluation:
+            return word_embeddings, pos_one_hots, caption, sent_len, motion, m_length, '_'.join(tokens)
+        else:
+            return caption, motion, m_length
+
+
+class Text2MotionDataset_Another_V(data.Dataset):
+    def __init__(self, mean, std, split_file, dataset_name, motion_dir, text_dir, unit_length, max_motion_length,
+                 max_text_length, evaluation=False, is_mesh=False):
+        self.evaluation = evaluation
+        self.max_length = 20
+        self.pointer = 0
+        self.max_motion_length = max_motion_length
+        self.max_text_len = max_text_length
+        self.unit_length = unit_length
+        min_motion_len = 40 if dataset_name =='t2m' else 24
+
+        data_dict = {}
+        id_list = []
+        with cs.open(split_file, 'r') as f:
+            for line in f.readlines():
+                id_list.append(line.strip())
+
+        new_name_list = []
+        length_list = []
+        for name in tqdm(id_list):
+            try:
+                motion = np.load(pjoin(motion_dir, name + '.npy'))
+                if not self.evaluation:
+                    if len(motion.shape) == 2:
+                        motion = np.expand_dims(motion, axis=0)
+                if is_mesh:
+                    if (len(motion)) < min_motion_len:
+                        continue
+                else:
+                    if (len(motion)) < min_motion_len or (len(motion) >= 200):
+                        continue
+                text_data = []
+                flag = False
+                with cs.open(pjoin(text_dir, name + '.txt')) as f:
+                    for line in f.readlines():
+                        text_dict = {}
+                        line_split = line.strip().split('#')
+                        caption = line_split[0]
+                        tokens = line_split[1].split(' ')
+                        f_tag = float(line_split[2])
+                        to_tag = float(line_split[3])
+                        f_tag = 0.0 if np.isnan(f_tag) else f_tag
+                        to_tag = 0.0 if np.isnan(to_tag) else to_tag
+
+                        text_dict['caption'] = caption
+                        text_dict['tokens'] = tokens
+                        if f_tag == 0.0 and to_tag == 0.0:
+                            flag = True
+                            text_data.append(text_dict)
+                        else:
+                            try:
+                                n_motion = motion[int(f_tag*20) : int(to_tag*20)]
+                                if (len(n_motion)) < min_motion_len or (len(n_motion) >= 200):
+                                    continue
+                                new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                                while new_name in data_dict:
+                                    new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                                data_dict[new_name] = {'motion': n_motion,
+                                                       'length': len(n_motion),
+                                                       'text':[text_dict]}
+                                new_name_list.append(new_name)
+                                length_list.append(len(n_motion))
+                            except:
+                                print(line_split)
+                                print(line_split[2], line_split[3], f_tag, to_tag, name)
+
+                if flag:
+                    data_dict[name] = {'motion': motion,
+                                       'length': len(motion),
+                                       'text': text_data}
+                    new_name_list.append(name)
+                    length_list.append(len(motion))
+            except:
+                pass
+        if self.evaluation:
+            self.w_vectorizer = GloVe('./glove', 'our_vab')
+            name_list, length_list = zip(*sorted(zip(new_name_list, length_list), key=lambda x: x[1]))
+        else:
+            name_list, length_list = new_name_list, length_list
+        self.mean = mean
+        self.std = std
+        self.length_arr = np.array(length_list)
+        self.data_dict = data_dict
+        self.name_list = name_list
+        if self.evaluation:
+            self.reset_max_len(self.max_length)
+
+    def reset_max_len(self, length):
+        assert length <= self.max_motion_length
+        self.pointer = np.searchsorted(self.length_arr, length)
+        print("Pointer Pointing at %d"%self.pointer)
+        self.max_length = length
+
+    def transform(self, data, mean=None, std=None):
+        if mean is None and std is None:
+            return (data - self.mean) / self.std
+        else:
+            return (data - mean) / std
+
+    def inv_transform(self, data, mean=None, std=None):
+        if mean is None and std is None:
+            return data * self.std + self.mean
+        else:
+            return data * std + mean
+
+    def __len__(self):
+        return len(self.data_dict) - self.pointer
+
+    def __getitem__(self, item):
+        idx = self.pointer + item
+        data = self.data_dict[self.name_list[idx]]
+        motion, m_length, text_list = data['motion'], data['length'], data['text']
+        # Randomly select a caption
+        text_data = random.choice(text_list)
+        caption, tokens = text_data['caption'], text_data['tokens']
+
+        if self.evaluation:
+            if len(tokens) < self.max_text_len:
+                # pad with "unk"
+                tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+                sent_len = len(tokens)
+                tokens = tokens + ['unk/OTHER'] * (self.max_text_len + 2 - sent_len)
+            else:
+                # crop
+                tokens = tokens[:self.max_text_len]
+                tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+                sent_len = len(tokens)
+            pos_one_hots = []
+            word_embeddings = []
+            for token in tokens:
+                word_emb, pos_oh = self.w_vectorizer[token]
+                pos_one_hots.append(pos_oh[None, :])
+                word_embeddings.append(word_emb[None, :])
+            pos_one_hots = np.concatenate(pos_one_hots, axis=0)
+            word_embeddings = np.concatenate(word_embeddings, axis=0)
+
+        if self.unit_length < 10:
+            coin2 = np.random.choice(['single', 'single', 'double'])
+        else:
+            coin2 = 'single'
+
+        if coin2 == 'double':
+            m_length = (m_length // self.unit_length - 1) * self.unit_length
+        elif coin2 == 'single':
+            m_length = (m_length // self.unit_length) * self.unit_length
+        idx = random.randint(0, len(motion) - m_length)
+        motion = motion[idx:idx+m_length]
+
+        "Z Normalization"
+        if self.evaluation:
+            motion = motion[:, :self.mean.shape[0]]
+        motion = (motion - self.mean) / self.std
+
+        if m_length < self.max_motion_length:
+            if self.evaluation:
+                motion = np.concatenate([motion,
+                                         np.zeros((self.max_motion_length - m_length, motion.shape[1]))
+                                         ], axis=0)
+            else:
+                motion = np.concatenate([motion,
+                                         np.zeros((self.max_motion_length - m_length, motion.shape[1], motion.shape[2]))
+                                         ], axis=0)
+        elif m_length > self.max_motion_length:
+            if not self.evaluation:
+                idx = random.randint(0, m_length - self.max_motion_length)
+                motion = motion[idx:idx + self.max_motion_length]
+        if self.evaluation:
+            return word_embeddings, pos_one_hots, caption, sent_len, motion, m_length, '_'.join(tokens)
+        else:
+            return caption, motion, m_length

utils/eval_mean_std/t2m/eval_mean.npy

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

utils/eval_mean_std/t2m/eval_std.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7898ab84360e67ef06046df9e39001b02182d3aeae2ae4e3a3f46277bd748045
+size 396

utils/eval_utils.py

@@ -0,0 +1,928 @@
+import os
+import numpy as np
+from scipy import linalg
+from scipy.ndimage import uniform_filter1d
+import torch
+from utils.motion_process import recover_from_ric
+from utils.back_process import back_process
+from tqdm import tqdm
+
+#################################################################################
+#                               Eval Function Loops                             #
+#################################################################################
+@torch.no_grad()
+def evaluation_ae(out_dir, val_loader, net, writer, ep, eval_wrapper, device, best_fid=1000, best_div=0,
+                  best_top1=0, best_top2=0, best_top3=0, best_matching=100,
+                  eval_mean=None, eval_std=None, save=True, draw=True):
+    net.eval()
+
+    motion_annotation_list = []
+    motion_pred_list = []
+
+    R_precision_real = 0
+    R_precision = 0
+
+    nb_sample = 0
+    matching_score_real = 0
+    matching_score_pred = 0
+    mpjpe = 0
+    num_poses = 0
+
+    for batch in tqdm(val_loader):
+        word_embeddings, pos_one_hots, caption, sent_len, motion, m_length, token = batch
+
+        bs, seq = motion.shape[0], motion.shape[1]
+        gt = val_loader.dataset.inv_transform(motion.detach().cpu().numpy())
+        bgt = []
+        for j in range(bs):
+            bgt.append(back_process(gt[j], is_mesh=False))
+        bgt = np.stack(bgt, axis=0)
+        bgt = val_loader.dataset.transform(bgt, eval_mean, eval_std)
+
+        bgt = torch.from_numpy(bgt).to(device)
+        (et, em), (_, _) = eval_wrapper.get_co_embeddings(word_embeddings, pos_one_hots, sent_len, caption, bgt, m_length-1)
+
+        motion = motion.float().to(device)
+        with torch.no_grad():
+            pred_pose_eval = net.forward(motion)
+        pred = val_loader.dataset.inv_transform(pred_pose_eval.detach().cpu().numpy())
+        bpred = []
+        for j in range(bs):
+            bpred.append(back_process(pred[j], is_mesh=False))
+        bpred = np.stack(bpred, axis=0)
+        bpred = val_loader.dataset.transform(bpred, eval_mean, eval_std)
+
+        (et_pred, em_pred), (_, _) = eval_wrapper.get_co_embeddings(word_embeddings, pos_one_hots, sent_len, caption,
+                                                          torch.from_numpy(bpred).to(device), m_length-1)
+        for i in range(bs):
+            gtt = torch.from_numpy(gt[i, :m_length[i]]).float().reshape(-1, 22, 3)
+            predd = torch.from_numpy(pred[i, :m_length[i]]).float().reshape(-1, 22, 3)
+            mpjpe += torch.sum(calculate_mpjpe(gtt, predd))
+            num_poses += gt.shape[0]
+
+        motion_pred_list.append(em_pred)
+        motion_annotation_list.append(em)
+
+        temp_R = calculate_R_precision(et.cpu().numpy(), em.cpu().numpy(), top_k=3, sum_all=True)
+        temp_match = euclidean_distance_matrix(et.cpu().numpy(), em.cpu().numpy()).trace()
+        R_precision_real += temp_R
+        matching_score_real += temp_match
+        temp_R = calculate_R_precision(et_pred.cpu().numpy(), em_pred.cpu().numpy(), top_k=3, sum_all=True)
+        temp_match = euclidean_distance_matrix(et_pred.cpu().numpy(), em_pred.cpu().numpy()).trace()
+        R_precision += temp_R
+        matching_score_pred += temp_match
+
+        nb_sample += bs
+
+    motion_annotation_np = torch.cat(motion_annotation_list, dim=0).cpu().numpy()
+    motion_pred_np = torch.cat(motion_pred_list, dim=0).cpu().numpy()
+    gt_mu, gt_cov = calculate_activation_statistics(motion_annotation_np)
+    mu, cov = calculate_activation_statistics(motion_pred_np)
+
+    diversity_real = calculate_diversity(motion_annotation_np, 300 if nb_sample > 300 else 100)
+    diversity = calculate_diversity(motion_pred_np, 300 if nb_sample > 300 else 100)
+
+    R_precision_real = R_precision_real / nb_sample
+    R_precision = R_precision / nb_sample
+
+    matching_score_real = matching_score_real / nb_sample
+    matching_score_pred = matching_score_pred / nb_sample
+    mpjpe = mpjpe / num_poses
+
+    fid = calculate_frechet_distance(gt_mu, gt_cov, mu, cov)
+
+    msg = "--> \t Eva. Re %d:, FID. %.4f, Diversity Real. %.4f, Diversity. %.4f, R_precision_real. (%.4f, %.4f, %.4f), R_precision. (%.4f, %.4f, %.4f), matching_real. %.4f, matching_pred. %.4f, MPJPE. %.4f" % \
+          (ep, fid, diversity_real, diversity, R_precision_real[0], R_precision_real[1], R_precision_real[2],
+           R_precision[0], R_precision[1], R_precision[2], matching_score_real, matching_score_pred, mpjpe)
+    print(msg)
+    if draw:
+        writer.add_scalar('./Test/FID', fid, ep)
+        writer.add_scalar('./Test/Diversity', diversity, ep)
+        writer.add_scalar('./Test/top1', R_precision[0], ep)
+        writer.add_scalar('./Test/top2', R_precision[1], ep)
+        writer.add_scalar('./Test/top3', R_precision[2], ep)
+        writer.add_scalar('./Test/matching_score', matching_score_pred, ep)
+
+    if fid < best_fid:
+        msg = "--> --> \t FID Improved from %.5f to %.5f !!!" % (best_fid, fid)
+        if draw: print(msg)
+        best_fid = fid
+        if save:
+            torch.save({'ae': net.state_dict(), 'ep': ep}, os.path.join(out_dir, 'net_best_fid.tar'))
+
+    if abs(diversity_real - diversity) < abs(diversity_real - best_div):
+        msg = "--> --> \t Diversity Improved from %.5f to %.5f !!!"%(best_div, diversity)
+        if draw: print(msg)
+        best_div = diversity
+
+    if R_precision[0] > best_top1:
+        msg = "--> --> \t Top1 Improved from %.5f to %.5f !!!" % (best_top1, R_precision[0])
+        if draw: print(msg)
+        best_top1 = R_precision[0]
+
+    if R_precision[1] > best_top2:
+        msg = "--> --> \t Top2 Improved from %.5f to %.5f!!!" % (best_top2, R_precision[1])
+        if draw: print(msg)
+        best_top2 = R_precision[1]
+
+    if R_precision[2] > best_top3:
+        msg = "--> --> \t Top3 Improved from %.5f to %.5f !!!" % (best_top3, R_precision[2])
+        if draw: print(msg)
+        best_top3 = R_precision[2]
+
+    if matching_score_pred < best_matching:
+        msg = f"--> --> \t matching_score Improved from %.5f to %.5f !!!" % (best_matching, matching_score_pred)
+        if draw: print(msg)
+        best_matching = matching_score_pred
+
+    net.train()
+    return best_fid, best_div, best_top1, best_top2, best_top3, best_matching, mpjpe, writer
+
+@torch.no_grad()
+def evaluation_acmdm(out_dir, val_loader, ema_acmdm, ae, writer, ep, best_fid, best_div,
+                        best_top1, best_top2, best_top3, best_matching, eval_wrapper, device, clip_score_old,
+                        cond_scale=None, cal_mm=False, eval_mean=None, eval_std=None, after_mean=None, after_std=None, mesh_mean=None, mesh_std=None,
+                        draw=True,
+                        is_raw=False,
+                        is_prefix=False,
+                        is_control=False, index=[0], intensity=100,
+                        is_mesh=False):
+
+    ema_acmdm.eval()
+    if not is_raw:
+        ae.eval()
+
+    save=False
+
+    motion_annotation_list = []
+    motion_pred_list = []
+    motion_multimodality = []
+    R_precision_real = 0
+    R_precision = 0
+    matching_score_real = 0
+    matching_score_pred = 0
+    multimodality = 0
+    if cond_scale is None:
+        if "kit" in out_dir:
+            cond_scale = 2.5
+        else:
+            cond_scale = 2.5
+    clip_score_real = 0
+    clip_score_gt = 0
+    skate_ratio_sum = 0
+    dist_sum = 0
+    traj_err = []
+
+    nb_sample = 0
+    if cal_mm:
+        num_mm_batch = 3
+    else:
+        num_mm_batch = 0
+
+    for i, batch in enumerate(tqdm(val_loader)):
+        word_embeddings, pos_one_hots, clip_text, sent_len, pose, m_length, token = batch
+        m_length = m_length.to(device)
+
+        bs, seq = pose.shape[:2]
+        if i < num_mm_batch:
+            motion_multimodality_batch = []
+            batch_clip_score_pred = 0
+            for _ in tqdm(range(30)):
+                pred_latents = ema_acmdm.generate(clip_text, m_length//4 if not is_raw else m_length, cond_scale)
+
+                if not is_raw:
+                    pred_latents = val_loader.dataset.inv_transform(pred_latents.permute(0, 2, 3, 1).detach().cpu().numpy(),
+                                                                    after_mean, after_std)
+                    pred_latents = torch.from_numpy(pred_latents).to(device)
+                    with torch.no_grad():
+                        pred_motions = ae.decode(pred_latents.permute(0,3,1,2))
+                else:
+                    pred_motions = pred_latents.permute(0, 2, 3, 1)
+                pred_motions = val_loader.dataset.inv_transform(pred_motions.detach().cpu().numpy())
+                pred_motionss = []
+                for j in range(bs):
+                    pred_motionss.append(back_process(pred_motions[j], is_mesh=is_mesh))
+                pred_motionss = np.stack(pred_motionss, axis=0)
+                pred_motions = val_loader.dataset.transform(pred_motionss, eval_mean, eval_std)
+                (et_pred, em_pred), (et_pred_clip, em_pred_clip) = eval_wrapper.get_co_embeddings(word_embeddings,
+                                                                                                  pos_one_hots,
+                                                                                                  sent_len,
+                                                                                                  clip_text,
+                                                                                                  torch.from_numpy(
+                                                                                                      pred_motions).to(
+                                                                                                      device),
+                                                                                                  m_length - 1)
+                motion_multimodality_batch.append(em_pred.unsqueeze(1))
+            motion_multimodality_batch = torch.cat(motion_multimodality_batch, dim=1) #(bs, 30, d)
+            motion_multimodality.append(motion_multimodality_batch)
+            for j in range(bs):
+                single_em = em_pred_clip[j]
+                single_et = et_pred_clip[j]
+                clip_score = (single_em @ single_et.T).item()
+                batch_clip_score_pred += clip_score
+            clip_score_real += batch_clip_score_pred
+        else:
+            if is_control:
+                pred_latents, mask_hint = ema_acmdm.generate_control(clip_text, m_length//4, pose.clone().float().to(device).permute(0,3,1,2), index, intensity,
+                                                                     cond_scale)
+                mask_hint = mask_hint.permute(0, 2, 3, 1).cpu().numpy()
+            elif is_prefix:
+                motion = pose.clone().float().to(device)
+                with torch.no_grad():
+                    motion = ae.encode(motion)
+                amean = torch.from_numpy(after_mean).to(device)
+                astd = torch.from_numpy(after_std).to(device)
+                motion = ((motion.permute(0,2,3,1)-amean)/astd).permute(0,3,1,2)
+                pred_latents = ema_acmdm.generate(clip_text, m_length // 4, cond_scale, motion[:, :, :5, :]) # 20+40 style
+            else:
+                pred_latents = ema_acmdm.generate(clip_text, m_length//4 if not (is_raw or is_mesh) else m_length, cond_scale, j=22 if not is_mesh else 28)
+            if not is_raw:
+                pred_latents = val_loader.dataset.inv_transform(pred_latents.permute(0,2,3,1).detach().cpu().numpy(), after_mean, after_std)
+                pred_latents = torch.from_numpy(pred_latents).to(device)
+                with torch.no_grad():
+                    if not is_mesh:
+                        pred_latents = pred_latents.permute(0, 3, 1, 2)
+                    pred_motions = ae.decode(pred_latents)
+            else:
+                pred_motions = pred_latents.permute(0, 2, 3, 1)
+            if not is_mesh:
+                pred_motions = val_loader.dataset.inv_transform(pred_motions.detach().cpu().numpy())
+            else:
+                pred_motions = val_loader.dataset.inv_transform(pred_motions.detach().cpu().numpy(), mesh_mean, mesh_std)
+                J_regressor = np.load('body_models/J_regressor.npy')
+                pred_motions = np.einsum('jk,btkc->btjc', J_regressor, pred_motions)[:, :, :22]
+            if is_control:
+                # foot skate
+                skate_ratio, skate_vel = calculate_skating_ratio(torch.from_numpy(pred_motions.transpose(0,2,3,1)))  # [batch_size]
+                skate_ratio_sum += skate_ratio.sum()
+                # control errors
+                hint = val_loader.dataset.inv_transform(pose.clone().detach().cpu().numpy())
+                hint = hint * mask_hint
+                for i, (mot, h, mask) in enumerate(zip(pred_motions, hint, mask_hint)):
+                    control_error = control_l2(np.expand_dims(mot, axis=0), np.expand_dims(h, axis=0),
+                                               np.expand_dims(mask, axis=0))
+                    mean_error = control_error.sum() / mask.sum()
+                    dist_sum += mean_error
+                    control_error = control_error.reshape(-1)
+                    mask = mask.reshape(-1)
+                    err_np = calculate_trajectory_error(control_error, mean_error, mask)
+                    traj_err.append(err_np)
+
+            pred_motionss = []
+            for j in range(bs):
+                pred_motionss.append(back_process(pred_motions[j], is_mesh=is_mesh))
+            pred_motionss = np.stack(pred_motionss, axis=0)
+            pred_motions = val_loader.dataset.transform(pred_motionss, eval_mean, eval_std)
+            (et_pred, em_pred), (et_pred_clip, em_pred_clip) = eval_wrapper.get_co_embeddings(word_embeddings,
+                                                                              pos_one_hots, sent_len,
+                                                                              clip_text,
+                                                                              torch.from_numpy(pred_motions).to(device),
+                                                                              m_length-1)
+            batch_clip_score_pred = 0
+            for j in range(bs):
+                single_em = em_pred_clip[j]
+                single_et = et_pred_clip[j]
+                clip_score = (single_em @ single_et.T).item()
+                batch_clip_score_pred += clip_score
+            clip_score_real += batch_clip_score_pred
+
+        pose = val_loader.dataset.inv_transform(pose.detach().cpu().numpy())
+        poses = []
+        for j in range(bs):
+            poses.append(back_process(pose[j], is_mesh=False))
+        poses = np.stack(poses, axis=0)
+        pose = val_loader.dataset.transform(poses, eval_mean, eval_std)
+        pose = torch.from_numpy(pose).cuda().float()
+        pose = pose.cuda().float()
+        (et, em), (et_clip, em_clip) = eval_wrapper.get_co_embeddings(word_embeddings, pos_one_hots, sent_len, clip_text,
+                                                          pose.clone(), m_length-1)
+        batch_clip_score = 0
+        for j in range(bs):
+            single_em = em_clip[j]
+            single_et = et_clip[j]
+            clip_score = (single_em @ single_et.T).item()
+            batch_clip_score += clip_score
+        clip_score_gt += batch_clip_score
+        motion_annotation_list.append(em)
+        motion_pred_list.append(em_pred)
+
+        temp_R = calculate_R_precision(et.cpu().numpy(), em.cpu().numpy(), top_k=3, sum_all=True)
+        temp_match = euclidean_distance_matrix(et.cpu().numpy(), em.cpu().numpy()).trace()
+        R_precision_real += temp_R
+        matching_score_real += temp_match
+        temp_R = calculate_R_precision(et_pred.cpu().numpy(), em_pred.cpu().numpy(), top_k=3, sum_all=True)
+        temp_match = euclidean_distance_matrix(et_pred.cpu().numpy(), em_pred.cpu().numpy()).trace()
+        R_precision += temp_R
+        matching_score_pred += temp_match
+
+        nb_sample += bs
+
+    motion_annotation_np = torch.cat(motion_annotation_list, dim=0).cpu().numpy()
+    motion_pred_np = torch.cat(motion_pred_list, dim=0).cpu().numpy()
+    gt_mu, gt_cov = calculate_activation_statistics(motion_annotation_np)
+    mu, cov = calculate_activation_statistics(motion_pred_np)
+
+    diversity_real = calculate_diversity(motion_annotation_np, 300 if nb_sample > 300 else 100)
+    diversity = calculate_diversity(motion_pred_np, 300 if nb_sample > 300 else 100)
+
+    R_precision_real = R_precision_real / nb_sample
+    R_precision = R_precision / nb_sample
+
+    clip_score_real = clip_score_real / nb_sample
+    clip_score_gt = clip_score_gt / nb_sample
+
+    matching_score_real = matching_score_real / nb_sample
+    matching_score_pred = matching_score_pred / nb_sample
+    if is_control:
+        # l2 dist
+        dist_mean = dist_sum / nb_sample
+
+        # Skating evaluation
+        skating_score = skate_ratio_sum / nb_sample
+
+        ### For trajecotry evaluation from GMD ###
+        traj_err = np.stack(traj_err).mean(0)
+
+    if cal_mm:
+        motion_multimodality = torch.cat(motion_multimodality, dim=0).cpu().numpy()
+        multimodality = calculate_multimodality(motion_multimodality, 10)
+
+    fid = calculate_frechet_distance(gt_mu, gt_cov, mu, cov)
+
+    msg = (f"--> \t Eva. Ep/Re {ep} :, FID. {fid:.4f}, Diversity Real. {diversity_real:.4f}, Diversity."
+           f" {diversity:.4f}, R_precision_real. {R_precision_real}, R_precision. {R_precision},"
+           f" matching_score_real. {matching_score_real}, matching_score_pred. {matching_score_pred}"
+           f" multimodality. {multimodality:.4f}, clip score. {clip_score_real}"
+        + (f" foot skating. {skating_score:.4f}, traj error. {traj_err[1].item():.4f}, pos error. {traj_err[3].item():.4f}, avg error. {traj_err[4].item():.4f}"
+           if is_control else ""))
+    print(msg)
+
+    if draw:
+        writer.add_scalar('./Test/FID', fid, ep)
+        writer.add_scalar('./Test/Diversity', diversity, ep)
+        writer.add_scalar('./Test/top1', R_precision[0], ep)
+        writer.add_scalar('./Test/top2', R_precision[1], ep)
+        writer.add_scalar('./Test/top3', R_precision[2], ep)
+        writer.add_scalar('./Test/matching_score', matching_score_pred, ep)
+        writer.add_scalar('./Test/clip_score', clip_score_real, ep)
+
+
+    if fid < best_fid:
+        msg = f"--> --> \t FID Improved from {best_fid:.5f} to {fid:.5f} !!!"
+        if draw: print(msg)
+        best_fid, best_ep = fid, ep
+        save=True
+
+
+    if matching_score_pred < best_matching:
+        msg = f"--> --> \t matching_score Improved from {best_matching:.5f} to {matching_score_pred:.5f} !!!"
+        if draw: print(msg)
+        best_matching = matching_score_pred
+
+    if abs(diversity_real - diversity) < abs(diversity_real - best_div):
+        msg = f"--> --> \t Diversity Improved from {best_div:.5f} to {diversity:.5f} !!!"
+        if draw: print(msg)
+        best_div = diversity
+
+    if R_precision[0] > best_top1:
+        msg = f"--> --> \t Top1 Improved from {best_top1:.4f} to {R_precision[0]:.4f} !!!"
+        if draw: print(msg)
+        best_top1 = R_precision[0]
+
+    if R_precision[1] > best_top2:
+        msg = f"--> --> \t Top2 Improved from {best_top2:.4f} to {R_precision[1]:.4f} !!!"
+        if draw: print(msg)
+        best_top2 = R_precision[1]
+
+    if R_precision[2] > best_top3:
+        msg = f"--> --> \t Top3 Improved from {best_top3:.4f} to {R_precision[2]:.4f} !!!"
+        if draw: print(msg)
+        best_top3 = R_precision[2]
+
+    if clip_score_real > clip_score_old:
+        msg = f"--> --> \t CLIP-score Improved from {clip_score_old:.4f} to {clip_score_real:.4f} !!!"
+        if draw: print(msg)
+        clip_score_old = clip_score_real
+
+    if cal_mm:
+        return best_fid, best_div, best_top1, best_top2, best_top3, best_matching, multimodality, clip_score_old, writer, save
+    else:
+        if is_control:
+            return best_fid, best_div, best_top1, best_top2, best_top3, best_matching, 0, clip_score_old, writer, save, dist_mean, skating_score, traj_err
+        else:
+            return best_fid, best_div, best_top1, best_top2, best_top3, best_matching, 0, clip_score_old, writer, save, None, None, None
+
+
+@torch.no_grad()
+def evaluation_acmdm_another_v(out_dir, val_loader, ema_acmdm, ae, writer, ep, best_fid, best_div,
+                        best_top1, best_top2, best_top3, best_matching, eval_wrapper, device, clip_score_old,
+                        cond_scale=None, cal_mm=False, train_mean=None, train_std=None, after_mean=None, after_std=None,
+                        draw=True,
+                        is_raw=False,
+                        is_prefix=False,
+                        is_control=False, index=[0], intensity=100,
+                        is_mesh=False):
+
+    ema_acmdm.eval()
+    if not is_raw:
+        ae.eval()
+
+    save=False
+
+    motion_annotation_list = []
+    motion_pred_list = []
+    motion_multimodality = []
+    R_precision_real = 0
+    R_precision = 0
+    matching_score_real = 0
+    matching_score_pred = 0
+    multimodality = 0
+    if cond_scale is None:
+        if "kit" in out_dir:
+            cond_scale = 2.5
+        else:
+            cond_scale = 2.5
+    clip_score_real = 0
+    clip_score_gt = 0
+    skate_ratio_sum = 0
+    dist_sum = 0
+    traj_err = []
+
+    nb_sample = 0
+    if cal_mm:
+        num_mm_batch = 3
+    else:
+        num_mm_batch = 0
+
+    for i, batch in enumerate(tqdm(val_loader)):
+        word_embeddings, pos_one_hots, clip_text, sent_len, pose, m_length, token = batch
+        m_length = m_length.to(device)
+
+        bs, seq = pose.shape[:2]
+        if i < num_mm_batch:
+            motion_multimodality_batch = []
+            batch_clip_score_pred = 0
+            for _ in tqdm(range(30)):
+                pred_latents = ema_acmdm.generate(clip_text, m_length//4 if not is_raw else m_length, cond_scale)
+
+                if not is_raw:
+                    pred_latents = val_loader.dataset.inv_transform(pred_latents.permute(0, 2, 3, 1).detach().cpu().numpy(),
+                                                                    after_mean, after_std)
+                    pred_latents = torch.from_numpy(pred_latents).to(device)
+                    with torch.no_grad():
+                        pred_motions = ae.decode(pred_latents.permute(0,3,1,2))
+                else:
+                    pred_motions = pred_latents.permute(0, 2, 3, 1)
+                pred_motions = val_loader.dataset.inv_transform(pred_motions.detach().cpu().numpy(), train_mean, train_std)
+                pred_motionss = []
+                for j in range(bs):
+                    pred_motionss.append(back_process(pred_motions[j], is_mesh=is_mesh))
+                pred_motionss = np.stack(pred_motionss, axis=0)
+                pred_motions = val_loader.dataset.transform(pred_motionss)
+                (et_pred, em_pred), (et_pred_clip, em_pred_clip) = eval_wrapper.get_co_embeddings(word_embeddings,
+                                                                                                  pos_one_hots,
+                                                                                                  sent_len,
+                                                                                                  clip_text,
+                                                                                                  torch.from_numpy(
+                                                                                                      pred_motions).to(
+                                                                                                      device),
+                                                                                                  m_length - 1)
+                motion_multimodality_batch.append(em_pred.unsqueeze(1))
+            motion_multimodality_batch = torch.cat(motion_multimodality_batch, dim=1) #(bs, 30, d)
+            motion_multimodality.append(motion_multimodality_batch)
+            for j in range(bs):
+                single_em = em_pred_clip[j]
+                single_et = et_pred_clip[j]
+                clip_score = (single_em @ single_et.T).item()
+                batch_clip_score_pred += clip_score
+            clip_score_real += batch_clip_score_pred
+        else:
+            if is_control:
+                bgt = val_loader.dataset.inv_transform(pose.clone())
+                motion_gt = []
+                for j in range(bs):
+                    motion_gt.append(recover_from_ric(bgt[j].float(), 22).numpy())
+                motion_gt = np.stack(motion_gt, axis=0)
+                motion = val_loader.dataset.transform(motion_gt, train_mean, train_std)
+                motion = torch.from_numpy(motion).float().to(device)
+                pred_latents, mask_hint = ema_acmdm.generate_control(clip_text, m_length//4, motion.clone().permute(0,3,1,2), index, intensity,
+                                                                     cond_scale)
+                mask_hint = mask_hint.permute(0, 2, 3, 1).cpu().numpy()
+            elif is_prefix:
+                bgt = val_loader.dataset.inv_transform(pose.clone())
+                motion_gt = []
+                for j in range(bs):
+                    motion_gt.append(recover_from_ric(bgt[j].float(), 22).numpy())
+                motion_gt = np.stack(motion_gt, axis=0)
+                motion = val_loader.dataset.transform(motion_gt, train_mean, train_std)
+                motion = torch.from_numpy(motion).float().to(device)
+                with torch.no_grad():
+                    motion = ae.encode(motion)
+                amean = torch.from_numpy(after_mean).to(device)
+                astd = torch.from_numpy(after_std).to(device)
+                motion = ((motion.permute(0,2,3,1)-amean)/astd).permute(0,3,1,2)
+                pred_latents = ema_acmdm.generate(clip_text, m_length // 4, cond_scale, motion[:, :, :5, :]) # 20+40 style
+            else:
+                pred_latents = ema_acmdm.generate(clip_text, m_length//4 if not is_raw else m_length, cond_scale)
+            if not is_raw:
+                pred_latents = val_loader.dataset.inv_transform(pred_latents.permute(0,2,3,1).detach().cpu().numpy(), after_mean, after_std)
+                pred_latents = torch.from_numpy(pred_latents).to(device)
+                with torch.no_grad():
+                    pred_motions = ae.decode(pred_latents.permute(0,3,1,2))
+            else:
+                pred_motions = pred_latents.permute(0, 2, 3, 1)
+            pred_motions = val_loader.dataset.inv_transform(pred_motions.detach().cpu().numpy(), train_mean, train_std)
+            if is_control:
+                # foot skate
+                skate_ratio, skate_vel = calculate_skating_ratio(torch.from_numpy(pred_motions.transpose(0,2,3,1)))  # [batch_size]
+                skate_ratio_sum += skate_ratio.sum()
+                # control errors
+                hint = motion_gt * mask_hint
+                for i, (mot, h, mask) in enumerate(zip(pred_motions, hint, mask_hint)):
+                    control_error = control_l2(np.expand_dims(mot, axis=0), np.expand_dims(h, axis=0),
+                                               np.expand_dims(mask, axis=0))
+                    mean_error = control_error.sum() / mask.sum()
+                    dist_sum += mean_error
+                    control_error = control_error.reshape(-1)
+                    mask = mask.reshape(-1)
+                    err_np = calculate_trajectory_error(control_error, mean_error, mask)
+                    traj_err.append(err_np)
+
+            pred_motionss = []
+            for j in range(bs):
+                pred_motionss.append(back_process(pred_motions[j], is_mesh=is_mesh))
+            pred_motionss = np.stack(pred_motionss, axis=0)
+            pred_motions = val_loader.dataset.transform(pred_motionss)
+            (et_pred, em_pred), (et_pred_clip, em_pred_clip) = eval_wrapper.get_co_embeddings(word_embeddings,
+                                                                              pos_one_hots, sent_len,
+                                                                              clip_text,
+                                                                              torch.from_numpy(pred_motions).to(device),
+                                                                              m_length-1)
+            batch_clip_score_pred = 0
+            for j in range(bs):
+                single_em = em_pred_clip[j]
+                single_et = et_pred_clip[j]
+                clip_score = (single_em @ single_et.T).item()
+                batch_clip_score_pred += clip_score
+            clip_score_real += batch_clip_score_pred
+
+        pose = pose.cuda().float()
+        (et, em), (et_clip, em_clip) = eval_wrapper.get_co_embeddings(word_embeddings, pos_one_hots, sent_len, clip_text,
+                                                          pose.clone(), m_length-1)
+        batch_clip_score = 0
+        for j in range(bs):
+            single_em = em_clip[j]
+            single_et = et_clip[j]
+            clip_score = (single_em @ single_et.T).item()
+            batch_clip_score += clip_score
+        clip_score_gt += batch_clip_score
+        motion_annotation_list.append(em)
+        motion_pred_list.append(em_pred)
+
+        temp_R = calculate_R_precision(et.cpu().numpy(), em.cpu().numpy(), top_k=3, sum_all=True)
+        temp_match = euclidean_distance_matrix(et.cpu().numpy(), em.cpu().numpy()).trace()
+        R_precision_real += temp_R
+        matching_score_real += temp_match
+        temp_R = calculate_R_precision(et_pred.cpu().numpy(), em_pred.cpu().numpy(), top_k=3, sum_all=True)
+        temp_match = euclidean_distance_matrix(et_pred.cpu().numpy(), em_pred.cpu().numpy()).trace()
+        R_precision += temp_R
+        matching_score_pred += temp_match
+
+        nb_sample += bs
+
+    motion_annotation_np = torch.cat(motion_annotation_list, dim=0).cpu().numpy()
+    motion_pred_np = torch.cat(motion_pred_list, dim=0).cpu().numpy()
+    gt_mu, gt_cov = calculate_activation_statistics(motion_annotation_np)
+    mu, cov = calculate_activation_statistics(motion_pred_np)
+
+    diversity_real = calculate_diversity(motion_annotation_np, 300 if nb_sample > 300 else 100)
+    diversity = calculate_diversity(motion_pred_np, 300 if nb_sample > 300 else 100)
+
+    R_precision_real = R_precision_real / nb_sample
+    R_precision = R_precision / nb_sample
+
+    clip_score_real = clip_score_real / nb_sample
+    clip_score_gt = clip_score_gt / nb_sample
+
+    matching_score_real = matching_score_real / nb_sample
+    matching_score_pred = matching_score_pred / nb_sample
+    if is_control:
+        # l2 dist
+        dist_mean = dist_sum / nb_sample
+
+        # Skating evaluation
+        skating_score = skate_ratio_sum / nb_sample
+
+        ### For trajecotry evaluation from GMD ###
+        traj_err = np.stack(traj_err).mean(0)
+
+    if cal_mm:
+        motion_multimodality = torch.cat(motion_multimodality, dim=0).cpu().numpy()
+        multimodality = calculate_multimodality(motion_multimodality, 10)
+
+    fid = calculate_frechet_distance(gt_mu, gt_cov, mu, cov)
+
+    msg = (f"--> \t Eva. Ep/Re {ep} :, FID. {fid:.4f}, Diversity Real. {diversity_real:.4f}, Diversity."
+           f" {diversity:.4f}, R_precision_real. {R_precision_real}, R_precision. {R_precision},"
+           f" matching_score_real. {matching_score_real}, matching_score_pred. {matching_score_pred}"
+           f" multimodality. {multimodality:.4f}, clip score. {clip_score_real}"
+        + (f" foot skating. {skating_score:.4f}, traj error. {traj_err[1].item():.4f}, loc error. {traj_err[3].item():.4f}, avg error. {traj_err[4].item():.4f}"
+           if is_control else ""))
+    print(msg)
+
+    if draw:
+        writer.add_scalar('./Test/FID', fid, ep)
+        writer.add_scalar('./Test/Diversity', diversity, ep)
+        writer.add_scalar('./Test/top1', R_precision[0], ep)
+        writer.add_scalar('./Test/top2', R_precision[1], ep)
+        writer.add_scalar('./Test/top3', R_precision[2], ep)
+        writer.add_scalar('./Test/matching_score', matching_score_pred, ep)
+        writer.add_scalar('./Test/clip_score', clip_score_real, ep)
+
+
+    if fid < best_fid:
+        msg = f"--> --> \t FID Improved from {best_fid:.5f} to {fid:.5f} !!!"
+        if draw: print(msg)
+        best_fid, best_ep = fid, ep
+        save=True
+
+
+    if matching_score_pred < best_matching:
+        msg = f"--> --> \t matching_score Improved from {best_matching:.5f} to {matching_score_pred:.5f} !!!"
+        if draw: print(msg)
+        best_matching = matching_score_pred
+
+    if abs(diversity_real - diversity) < abs(diversity_real - best_div):
+        msg = f"--> --> \t Diversity Improved from {best_div:.5f} to {diversity:.5f} !!!"
+        if draw: print(msg)
+        best_div = diversity
+
+    if R_precision[0] > best_top1:
+        msg = f"--> --> \t Top1 Improved from {best_top1:.4f} to {R_precision[0]:.4f} !!!"
+        if draw: print(msg)
+        best_top1 = R_precision[0]
+
+    if R_precision[1] > best_top2:
+        msg = f"--> --> \t Top2 Improved from {best_top2:.4f} to {R_precision[1]:.4f} !!!"
+        if draw: print(msg)
+        best_top2 = R_precision[1]
+
+    if R_precision[2] > best_top3:
+        msg = f"--> --> \t Top3 Improved from {best_top3:.4f} to {R_precision[2]:.4f} !!!"
+        if draw: print(msg)
+        best_top3 = R_precision[2]
+
+    if clip_score_real > clip_score_old:
+        msg = f"--> --> \t CLIP-score Improved from {clip_score_old:.4f} to {clip_score_real:.4f} !!!"
+        if draw: print(msg)
+        clip_score_old = clip_score_real
+
+    if cal_mm:
+        return best_fid, best_div, best_top1, best_top2, best_top3, best_matching, multimodality, clip_score_old, writer, save
+    else:
+        if is_control:
+            return best_fid, best_div, best_top1, best_top2, best_top3, best_matching, 0, clip_score_old, writer, save, dist_mean, skating_score, traj_err
+        else:
+            return best_fid, best_div, best_top1, best_top2, best_top3, best_matching, 0, clip_score_old, writer, save, None, None, None
+
+#################################################################################
+#                                 Util Functions                                #
+#################################################################################
+def eval_decorator(fn):
+    def inner(model, *args, **kwargs):
+        was_training = model.training
+        model.eval()
+        out = fn(model, *args, **kwargs)
+        model.train(was_training)
+        return out
+    return inner
+
+#################################################################################
+#                                     Metrics                                   #
+#################################################################################
+def calculate_mpjpe(gt_joints, pred_joints):
+    """
+    gt_joints: num_poses x num_joints(22) x 3
+    pred_joints: num_poses x num_joints(22) x 3
+    (obtained from recover_from_ric())
+    """
+    assert gt_joints.shape == pred_joints.shape, f"GT shape: {gt_joints.shape}, pred shape: {pred_joints.shape}"
+
+    # Align by root (pelvis)
+    pelvis = gt_joints[:, [0]].mean(1)
+    gt_joints = gt_joints - torch.unsqueeze(pelvis, dim=1)
+    pelvis = pred_joints[:, [0]].mean(1)
+    pred_joints = pred_joints - torch.unsqueeze(pelvis, dim=1)
+
+    # Compute MPJPE
+    mpjpe = torch.linalg.norm(pred_joints - gt_joints, dim=-1) # num_poses x num_joints=22
+    mpjpe_seq = mpjpe.mean(-1) # num_poses
+
+    return mpjpe_seq
+
+# (X - X_train)*(X - X_train) = -2X*X_train + X*X + X_train*X_train
+def euclidean_distance_matrix(matrix1, matrix2):
+    """
+        Params:
+        -- matrix1: N1 x D
+        -- matrix2: N2 x D
+        Returns:
+        -- dist: N1 x N2
+        dist[i, j] == distance(matrix1[i], matrix2[j])
+    """
+    assert matrix1.shape[1] == matrix2.shape[1]
+    d1 = -2 * np.dot(matrix1, matrix2.T)    # shape (num_test, num_train)
+    d2 = np.sum(np.square(matrix1), axis=1, keepdims=True)    # shape (num_test, 1)
+    d3 = np.sum(np.square(matrix2), axis=1)     # shape (num_train, )
+    dists = np.sqrt(d1 + d2 + d3)  # broadcasting
+    return dists
+
+def calculate_top_k(mat, top_k):
+    size = mat.shape[0]
+    gt_mat = np.expand_dims(np.arange(size), 1).repeat(size, 1)
+    bool_mat = (mat == gt_mat)
+    correct_vec = False
+    top_k_list = []
+    for i in range(top_k):
+#         print(correct_vec, bool_mat[:, i])
+        correct_vec = (correct_vec | bool_mat[:, i])
+        # print(correct_vec)
+        top_k_list.append(correct_vec[:, None])
+    top_k_mat = np.concatenate(top_k_list, axis=1)
+    return top_k_mat
+
+
+def calculate_R_precision(embedding1, embedding2, top_k, sum_all=False):
+    dist_mat = euclidean_distance_matrix(embedding1, embedding2)
+    argmax = np.argsort(dist_mat, axis=1)
+    top_k_mat = calculate_top_k(argmax, top_k)
+    if sum_all:
+        return top_k_mat.sum(axis=0)
+    else:
+        return top_k_mat
+
+
+def calculate_matching_score(embedding1, embedding2, sum_all=False):
+    assert len(embedding1.shape) == 2
+    assert embedding1.shape[0] == embedding2.shape[0]
+    assert embedding1.shape[1] == embedding2.shape[1]
+
+    dist = linalg.norm(embedding1 - embedding2, axis=1)
+    if sum_all:
+        return dist.sum(axis=0)
+    else:
+        return dist
+
+
+
+def calculate_activation_statistics(activations):
+    """
+    Params:
+    -- activation: num_samples x dim_feat
+    Returns:
+    -- mu: dim_feat
+    -- sigma: dim_feat x dim_feat
+    """
+    mu = np.mean(activations, axis=0)
+    cov = np.cov(activations, rowvar=False)
+    return mu, cov
+
+
+def calculate_diversity(activation, diversity_times):
+    assert len(activation.shape) == 2
+    assert activation.shape[0] > diversity_times
+    num_samples = activation.shape[0]
+
+    first_indices = np.random.choice(num_samples, diversity_times, replace=False)
+    second_indices = np.random.choice(num_samples, diversity_times, replace=False)
+    dist = linalg.norm(activation[first_indices] - activation[second_indices], axis=1)
+    return dist.mean()
+
+
+def calculate_multimodality(activation, multimodality_times):
+    assert len(activation.shape) == 3
+    assert activation.shape[1] > multimodality_times
+    num_per_sent = activation.shape[1]
+
+    first_dices = np.random.choice(num_per_sent, multimodality_times, replace=False)
+    second_dices = np.random.choice(num_per_sent, multimodality_times, replace=False)
+    dist = linalg.norm(activation[:, first_dices] - activation[:, second_dices], axis=2)
+    return dist.mean()
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    """Numpy implementation of the Frechet Distance.
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+    Stable version by Dougal J. Sutherland.
+    Params:
+    -- mu1   : Numpy array containing the activations of a layer of the
+               inception net (like returned by the function 'get_predictions')
+               for generated samples.
+    -- mu2   : The sample mean over activations, precalculated on an
+               representative data set.
+    -- sigma1: The covariance matrix over activations for generated samples.
+    -- sigma2: The covariance matrix over activations, precalculated on an
+               representative data set.
+    Returns:
+    --   : The Frechet Distance.
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert mu1.shape == mu2.shape, \
+        'Training and test mean vectors have different lengths'
+    assert sigma1.shape == sigma2.shape, \
+        'Training and test covariances have different dimensions'
+
+    diff = mu1 - mu2
+
+    # Product might be almost singular
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = ('fid calculation produces singular product; '
+               'adding %s to diagonal of cov estimates') % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # Numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError('Imaginary component {}'.format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return (diff.dot(diff) + np.trace(sigma1) +
+            np.trace(sigma2) - 2 * tr_covmean)
+
+
+# directly from omnicontrol
+def calculate_skating_ratio(motions):
+    thresh_height = 0.05  # 10
+    fps = 20.0
+    thresh_vel = 0.50  # 20 cm /s
+    avg_window = 5  # frames
+
+    batch_size = motions.shape[0]
+    # 10 left, 11 right foot. XZ plane, y up
+    # motions [bs, 22, 3, max_len]
+    verts_feet = motions[:, [10, 11], :, :].detach().cpu().numpy()  # [bs, 2, 3, max_len]
+    verts_feet_plane_vel = np.linalg.norm(verts_feet[:, :, [0, 2], 1:] - verts_feet[:, :, [0, 2], :-1],
+                                          axis=2) * fps  # [bs, 2, max_len-1]
+    # [bs, 2, max_len-1]
+    vel_avg = uniform_filter1d(verts_feet_plane_vel, axis=-1, size=avg_window, mode='constant', origin=0)
+
+    verts_feet_height = verts_feet[:, :, 1, :]  # [bs, 2, max_len]
+    # If feet touch ground in agjecent frames
+    feet_contact = np.logical_and((verts_feet_height[:, :, :-1] < thresh_height),
+                                  (verts_feet_height[:, :, 1:] < thresh_height))  # [bs, 2, max_len - 1]
+    # skate velocity
+    skate_vel = feet_contact * vel_avg
+
+    # it must both skating in the current frame
+    skating = np.logical_and(feet_contact, (verts_feet_plane_vel > thresh_vel))
+    # and also skate in the windows of frames
+    skating = np.logical_and(skating, (vel_avg > thresh_vel))
+
+    # Both feet slide
+    skating = np.logical_or(skating[:, 0, :], skating[:, 1, :])  # [bs, max_len -1]
+    skating_ratio = np.sum(skating, axis=1) / skating.shape[1]
+
+    return skating_ratio, skate_vel
+
+# directly from omnicontrol
+def control_l2(motion, hint, hint_mask):
+    # motion: b, seq, 22, 3
+    # hint: b, seq, 22, 1
+    loss = np.linalg.norm((motion - hint) * hint_mask, axis=-1)
+    return loss
+
+# directly from omnicontrol
+def calculate_trajectory_error(dist_error, mean_err_traj, mask, strict=True):
+    ''' dist_error shape [5]: error for each kps in metre
+      Two threshold: 20 cm and 50 cm.
+    If mean error in sequence is more then the threshold, fails
+    return: traj_fail(0.2), traj_fail(0.5), all_kps_fail(0.2), all_kps_fail(0.5), all_mean_err.
+        Every metrics are already averaged.
+    '''
+    # mean_err_traj = dist_error.mean(1)
+    if strict:
+        # Traj fails if any of the key frame fails
+        traj_fail_02 = 1.0 - (dist_error <= 0.2).all()
+        traj_fail_05 = 1.0 - (dist_error <= 0.5).all()
+    else:
+        # Traj fails if the mean error of all keyframes more than the threshold
+        traj_fail_02 = (mean_err_traj > 0.2)
+        traj_fail_05 = (mean_err_traj > 0.5)
+    all_fail_02 = (dist_error > 0.2).sum() / mask.sum()
+    all_fail_05 = (dist_error > 0.5).sum() / mask.sum()
+
+    return np.array([traj_fail_02, traj_fail_05, all_fail_02, all_fail_05, dist_error.sum() / mask.sum()])

utils/evaluators.py

@@ -0,0 +1,392 @@
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+import math
+from torch.nn.utils.rnn import pack_padded_sequence
+import torch.nn.functional as F
+from utils.glove import POS_enumerator
+import clip
+
+#################################################################################
+#                                    Evaluators                                 #
+#################################################################################
+def build_evaluators(dim_pose, dataset_name, dim_movement_enc_hidden, dim_movement_latent, dim_word, dim_pos_ohot, dim_text_hidden,
+                     dim_coemb_hidden, dim_motion_hidden, checkpoints_dir, device):
+    movement_enc = MovementConvEncoder(dim_pose, dim_movement_enc_hidden, dim_movement_latent)
+    text_enc = TextEncoderBiGRUCo(word_size=dim_word,
+                                  pos_size=dim_pos_ohot,
+                                  hidden_size=dim_text_hidden,
+                                  output_size=dim_coemb_hidden,
+                                  device=device)
+
+    motion_enc = MotionEncoderBiGRUCo(input_size=dim_movement_latent,
+                                      hidden_size=dim_motion_hidden,
+                                      output_size=dim_coemb_hidden,
+                                      device=device)
+    contrast_model = MotionCLIP(dim_pose)
+
+    checkpoint = torch.load(os.path.join(checkpoints_dir, dataset_name, 'text_mot_match', 'model', 'finest.tar'),
+                            map_location=device)
+    checkpoint_clip = torch.load(os.path.join(checkpoints_dir, dataset_name, 'text_mot_match_clip', 'model', 'finest.tar'),
+                            map_location=device)
+    movement_enc.load_state_dict(checkpoint['movement_encoder'])
+    text_enc.load_state_dict(checkpoint['text_encoder'])
+    motion_enc.load_state_dict(checkpoint['motion_encoder'])
+    contrast_model.load_state_dict(checkpoint_clip['contrast_model'])
+    print('Loading Evaluators')
+    return text_enc, motion_enc, movement_enc, contrast_model
+
+class Evaluators(object):
+
+    def __init__(self, dataset_name, device):
+        if dataset_name == 't2m':
+            dim_pose = 67
+        elif dataset_name == 'kit':
+            dim_pose = 64
+        else:
+            raise KeyError('Dataset not Recognized!!!')
+
+        dim_word = 300
+        dim_pos_ohot = len(POS_enumerator)
+        dim_motion_hidden = 1024
+        dim_movement_enc_hidden = 512
+        dim_movement_latent = 512
+        dim_text_hidden = 512
+        dim_coemb_hidden = 512
+        checkpoints_dir = 'checkpoints'
+        self.unit_length=4
+
+        self.text_encoder, self.motion_encoder, self.movement_encoder, self.contrast_model \
+        = build_evaluators(dim_pose, dataset_name, dim_movement_enc_hidden, dim_movement_latent, dim_word,
+                            dim_pos_ohot, dim_text_hidden, dim_coemb_hidden, dim_motion_hidden, checkpoints_dir, device)
+        self.device = device
+
+        self.text_encoder.to(device)
+        self.motion_encoder.to(device)
+        self.movement_encoder.to(device)
+        self.contrast_model.to(device)
+
+        self.text_encoder.eval()
+        self.motion_encoder.eval()
+        self.movement_encoder.eval()
+        self.contrast_model.eval()
+
+    def get_co_embeddings(self, word_embs, pos_ohot, cap_lens, captions, motions, m_lens):
+        with torch.no_grad():
+            word_embs = word_embs.detach().to(self.device).float()
+            pos_ohot = pos_ohot.detach().to(self.device).float()
+            motions = motions.detach().to(self.device).float()
+
+            '''clip based'''
+            clip_em = self.contrast_model.encode_motion(motions.clone(), m_lens)
+            clip_et = self.contrast_model.encode_text(captions)
+            clip_em = clip_em / clip_em.norm(dim=1, keepdim=True)
+            clip_et = clip_et / clip_et.norm(dim=1, keepdim=True)
+
+            '''original architecture'''
+            align_idx = np.argsort(m_lens.data.tolist())[::-1].copy()
+            motions = motions[align_idx]
+            m_lens = m_lens[align_idx]
+
+            movements = self.movement_encoder(motions).detach()
+            m_lens = m_lens // self.unit_length
+            motion_embedding = self.motion_encoder(movements, m_lens)
+
+            text_embedding = self.text_encoder(word_embs, pos_ohot, cap_lens)
+            text_embedding = text_embedding[align_idx]
+        return (text_embedding, motion_embedding), (clip_et, clip_em)
+
+    def get_motion_embeddings(self, motions, m_lens):
+        with torch.no_grad():
+            motions = motions.detach().to(self.device).float()
+            '''clip based'''
+            clip_em = self.contrast_model.encode_motion(motions.clone(), m_lens)
+            clip_em = clip_em / clip_em.norm(dim=1, keepdim=True)
+
+            '''original architecture'''
+            align_idx = np.argsort(m_lens.data.tolist())[::-1].copy()
+            motions = motions[align_idx]
+            m_lens = m_lens[align_idx]
+
+            movements = self.movement_encoder(motions).detach()
+            m_lens = m_lens // self.unit_length
+            motion_embedding = self.motion_encoder(movements, m_lens)
+        return motion_embedding, clip_em
+
+#################################################################################
+#                                 Inner Architectures                           #
+#################################################################################
+def init_weight(m):
+    if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear) or isinstance(m, nn.ConvTranspose1d):
+        nn.init.xavier_normal_(m.weight)
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=300):
+        super(PositionalEncoding, self).__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, pos):
+        return self.pe[pos]
+
+
+class PositionalEncodingCLIP(nn.Module):
+    def __init__(self, d_model, dropout=0.0, max_len=5000):
+        super(PositionalEncodingCLIP, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + self.pe[:x.shape[1], :].unsqueeze(0)
+        return self.dropout(x)
+
+
+def no_grad(nets):
+    if not isinstance(nets, list):
+        nets = [nets]
+    for net in nets:
+        if net is not None:
+            for param in net.parameters():
+                param.requires_grad = False
+
+def lengths_to_mask(lengths, max_len):
+    mask = torch.arange(max_len, device=lengths.device).expand(len(lengths), max_len) < lengths.unsqueeze(1)
+    return mask
+
+
+class MovementConvEncoder(nn.Module):
+    def __init__(self, input_size, hidden_size, output_size):
+        super(MovementConvEncoder, self).__init__()
+        self.main = nn.Sequential(
+            nn.Conv1d(input_size, hidden_size, 4, 2, 1),
+            nn.Dropout(0.2, inplace=True),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv1d(hidden_size, output_size, 4, 2, 1),
+            nn.Dropout(0.2, inplace=True),
+            nn.LeakyReLU(0.2, inplace=True),
+        )
+        self.out_net = nn.Linear(output_size, output_size)
+        self.main.apply(init_weight)
+        self.out_net.apply(init_weight)
+
+    def forward(self, inputs):
+        inputs = inputs.permute(0, 2, 1)
+        outputs = self.main(inputs).permute(0, 2, 1)
+        return self.out_net(outputs)
+
+
+class MovementConvDecoder(nn.Module):
+    def __init__(self, input_size, hidden_size, output_size):
+        super(MovementConvDecoder, self).__init__()
+        self.main = nn.Sequential(
+            nn.ConvTranspose1d(input_size, hidden_size, 4, 2, 1),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.ConvTranspose1d(hidden_size, output_size, 4, 2, 1),
+            nn.LeakyReLU(0.2, inplace=True),
+        )
+        self.out_net = nn.Linear(output_size, output_size)
+
+        self.main.apply(init_weight)
+        self.out_net.apply(init_weight)
+
+    def forward(self, inputs):
+        inputs = inputs.permute(0, 2, 1)
+        outputs = self.main(inputs).permute(0, 2, 1)
+        return self.out_net(outputs)
+
+class TextEncoderBiGRUCo(nn.Module):
+    def __init__(self, word_size, pos_size, hidden_size, output_size, device):
+        super(TextEncoderBiGRUCo, self).__init__()
+        self.device = device
+
+        self.pos_emb = nn.Linear(pos_size, word_size)
+        self.input_emb = nn.Linear(word_size, hidden_size)
+        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True)
+        self.output_net = nn.Sequential(
+            nn.Linear(hidden_size * 2, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(hidden_size, output_size)
+        )
+
+        self.input_emb.apply(init_weight)
+        self.pos_emb.apply(init_weight)
+        self.output_net.apply(init_weight)
+        self.hidden_size = hidden_size
+        self.hidden = nn.Parameter(torch.randn((2, 1, self.hidden_size), requires_grad=True))
+
+    def forward(self, word_embs, pos_onehot, cap_lens):
+        num_samples = word_embs.shape[0]
+
+        pos_embs = self.pos_emb(pos_onehot)
+        inputs = word_embs + pos_embs
+        input_embs = self.input_emb(inputs)
+        hidden = self.hidden.repeat(1, num_samples, 1)
+
+        cap_lens = cap_lens.data.tolist()
+        emb = pack_padded_sequence(input_embs, cap_lens, batch_first=True)
+
+        gru_seq, gru_last = self.gru(emb, hidden)
+
+        gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1)
+
+        return self.output_net(gru_last)
+
+
+class MotionEncoderBiGRUCo(nn.Module):
+    def __init__(self, input_size, hidden_size, output_size, device):
+        super(MotionEncoderBiGRUCo, self).__init__()
+        self.device = device
+
+        self.input_emb = nn.Linear(input_size, hidden_size)
+        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True)
+        self.output_net = nn.Sequential(
+            nn.Linear(hidden_size*2, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(hidden_size, output_size)
+        )
+
+        self.input_emb.apply(init_weight)
+        self.output_net.apply(init_weight)
+        self.hidden_size = hidden_size
+        self.hidden = nn.Parameter(torch.randn((2, 1, self.hidden_size), requires_grad=True))
+
+    def forward(self, inputs, m_lens):
+        num_samples = inputs.shape[0]
+
+        input_embs = self.input_emb(inputs)
+        hidden = self.hidden.repeat(1, num_samples, 1)
+
+        cap_lens = m_lens.data.tolist()
+        emb = pack_padded_sequence(input_embs, cap_lens, batch_first=True)
+
+        gru_seq, gru_last = self.gru(emb, hidden)
+
+        gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1)
+
+        return self.output_net(gru_last)
+
+
+class MotionEncoder(nn.Module):
+    def __init__(self, in_dim, latent_dim, ff_size, num_layers, num_heads, dropout, activation):
+        super().__init__()
+        self.input_feats = in_dim
+        self.latent_dim = latent_dim
+        self.ff_size = ff_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.activation = activation
+
+        self.query_token = nn.Parameter(torch.randn(1, self.latent_dim))
+
+        self.embed_motion = nn.Linear(self.input_feats, self.latent_dim)
+        self.sequence_pos_encoder = PositionalEncodingCLIP(self.latent_dim, self.dropout, max_len=2000)
+
+        seqTransEncoderLayer = nn.TransformerEncoderLayer(d_model=self.latent_dim,
+                                                          nhead=self.num_heads,
+                                                          dim_feedforward=self.ff_size,
+                                                          dropout=self.dropout,
+                                                          activation=self.activation,)
+        self.transformer = nn.TransformerEncoder(seqTransEncoderLayer, num_layers=self.num_layers)
+        self.out_ln = nn.LayerNorm(self.latent_dim)
+        self.out = nn.Linear(self.latent_dim, 512)
+
+
+    def forward(self, motion, padding_mask):
+        B, T, D  = motion.shape
+
+        x_emb = self.embed_motion(motion)
+
+        emb = torch.cat([self.query_token[torch.zeros(B, dtype=torch.long, device=motion.device)][:,None], x_emb], dim=1)
+
+        padding_mask = torch.cat([torch.zeros_like(padding_mask[:, 0:1]), padding_mask], dim=1)
+
+        h = self.sequence_pos_encoder(emb)
+        h = h.permute(1, 0, 2)
+        h = self.transformer(h, src_key_padding_mask=padding_mask)
+        h = h.permute(1, 0, 2)
+        h = self.out_ln(h)
+        motion_emb = self.out(h[:,0])
+
+        return motion_emb
+
+
+class MotionCLIP(nn.Module):
+    def __init__(self, in_dim):
+        super().__init__()
+        self.motion_encoder = MotionEncoder(in_dim, 512, 1024, 8, 8, 0.2, 'gelu')
+        clip_model, _ = clip.load("ViT-B/16", device="cpu", jit=False)
+        self.token_embedding = clip_model.token_embedding
+        self.positional_embedding = clip_model.positional_embedding
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        no_grad(self.token_embedding)
+
+        textTransEncoderLayer = nn.TransformerEncoderLayer(
+            d_model=512,
+            nhead=8,
+            dim_feedforward=1024,
+            dropout=0.2,
+            activation="gelu",)
+        self.textTransEncoder = nn.TransformerEncoder(
+            textTransEncoderLayer,
+            num_layers=8)
+        self.text_ln = nn.LayerNorm(512)
+        self.out = nn.Linear(512, 512)
+
+    def encode_motion(self, motion, m_lens):
+        seq_len = motion.shape[1]
+        padding_mask = ~lengths_to_mask(m_lens, seq_len)
+        motion_embedding = self.motion_encoder(motion, padding_mask.to(motion.device))
+        return motion_embedding
+
+    def encode_text(self, text):
+        device = next(self.parameters()).device
+
+        with torch.no_grad():
+            text = clip.tokenize(text, truncate=True).to(device)
+            x = self.token_embedding(text).float()
+            pe_tokens = x + self.positional_embedding.float()
+        pe_tokens = pe_tokens.permute(1,0,2)
+        out = self.textTransEncoder(pe_tokens)
+        out = out.permute(1, 0, 2)
+        out = self.text_ln(out)
+
+        out = out[torch.arange(x.shape[0]), text.argmax(dim=-1)]
+        out = self.out(out)
+        return out
+
+    def forward(self, motion, m_lens, text):
+        motion_features = self.encode_motion(motion, m_lens)
+        text_features = self.encode_text(text)
+
+        motion_features = motion_features / motion_features .norm(dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+        logit_scale = self.logit_scale.exp()
+        logits_per_motion = logit_scale * motion_features @ text_features.t()
+        logits_per_text = logits_per_motion.t()
+        return logits_per_motion, logits_per_text
+
+    def forward_loss(self, motion, m_lens, text):
+        logits_per_motion, logits_per_text = self.forward(motion, m_lens, text)
+        labels = torch.arange(len(logits_per_motion)).to(logits_per_motion.device)
+
+        image_loss = F.cross_entropy(logits_per_motion, labels)
+        text_loss = F.cross_entropy(logits_per_text, labels)
+        loss = (image_loss + text_loss) / 2
+        return loss

utils/glove.py

@@ -0,0 +1,83 @@
+import numpy as np
+import pickle
+from os.path import join as pjoin
+
+#################################################################################
+#                                      GloVe                                    #
+#################################################################################
+POS_enumerator = {
+    'VERB': 0,
+    'NOUN': 1,
+    'DET': 2,
+    'ADP': 3,
+    'NUM': 4,
+    'AUX': 5,
+    'PRON': 6,
+    'ADJ': 7,
+    'ADV': 8,
+    'Loc_VIP': 9,
+    'Body_VIP': 10,
+    'Obj_VIP': 11,
+    'Act_VIP': 12,
+    'Desc_VIP': 13,
+    'OTHER': 14,
+}
+
+Loc_list = ('left', 'right', 'clockwise', 'counterclockwise', 'anticlockwise', 'forward', 'back', 'backward',
+            'up', 'down', 'straight', 'curve')
+
+Body_list = ('arm', 'chin', 'foot', 'feet', 'face', 'hand', 'mouth', 'leg', 'waist', 'eye', 'knee', 'shoulder', 'thigh')
+
+Obj_List = ('stair', 'dumbbell', 'chair', 'window', 'floor', 'car', 'ball', 'handrail', 'baseball', 'basketball')
+
+Act_list = ('walk', 'run', 'swing', 'pick', 'bring', 'kick', 'put', 'squat', 'throw', 'hop', 'dance', 'jump', 'turn',
+            'stumble', 'dance', 'stop', 'sit', 'lift', 'lower', 'raise', 'wash', 'stand', 'kneel', 'stroll',
+            'rub', 'bend', 'balance', 'flap', 'jog', 'shuffle', 'lean', 'rotate', 'spin', 'spread', 'climb')
+
+Desc_list = ('slowly', 'carefully', 'fast', 'careful', 'slow', 'quickly', 'happy', 'angry', 'sad', 'happily',
+             'angrily', 'sadly')
+
+VIP_dict = {
+    'Loc_VIP': Loc_list,
+    'Body_VIP': Body_list,
+    'Obj_VIP': Obj_List,
+    'Act_VIP': Act_list,
+    'Desc_VIP': Desc_list,
+}
+
+
+class GloVe(object):
+    def __init__(self, meta_root, prefix):
+        vectors = np.load(pjoin(meta_root, '%s_data.npy'%prefix))
+        words = pickle.load(open(pjoin(meta_root, '%s_words.pkl'%prefix), 'rb'))
+        self.word2idx = pickle.load(open(pjoin(meta_root, '%s_idx.pkl'%prefix), 'rb'))
+        self.word2vec = {w: vectors[self.word2idx[w]] for w in words}
+
+    def _get_pos_ohot(self, pos):
+        pos_vec = np.zeros(len(POS_enumerator))
+        if pos in POS_enumerator:
+            pos_vec[POS_enumerator[pos]] = 1
+        else:
+            pos_vec[POS_enumerator['OTHER']] = 1
+        return pos_vec
+
+    def __len__(self):
+        return len(self.word2vec)
+
+    def __getitem__(self, item):
+        word, pos = item.split('/')
+        if word in self.word2vec:
+            word_vec = self.word2vec[word]
+            vip_pos = None
+            for key, values in VIP_dict.items():
+                if word in values:
+                    vip_pos = key
+                    break
+            if vip_pos is not None:
+                pos_vec = self._get_pos_ohot(vip_pos)
+            else:
+                pos_vec = self._get_pos_ohot(pos)
+        else:
+            word_vec = self.word2vec['unk']
+            pos_vec = self._get_pos_ohot('OTHER')
+        return word_vec, pos_vec

utils/mesh_mean_std/t2m/mesh_mean.npy

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

utils/mesh_mean_std/t2m/mesh_std.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19f56c7aaa752639b39eb76742b5ee1b79fb71931b8412dd84ca7312398293b4
+size 140

utils/motion_process.py

@@ -0,0 +1,429 @@
+import torch
+import numpy as np
+import matplotlib
+import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation, writers
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+import mpl_toolkits.mplot3d.axes3d as p3
+import os
+import io
+try:
+    from PIL import Image
+except ImportError:
+    Image = None
+try:
+    import imageio
+except ImportError:
+    imageio = None
+
+#################################################################################
+#                                   Data Params                                 #
+#################################################################################
+kit_kinematic_chain = [[0, 11, 12, 13, 14, 15], [0, 16, 17, 18, 19, 20], [0, 1, 2, 3, 4], [3, 5, 6, 7], [3, 8, 9, 10]]
+t2m_kinematic_chain = [[0, 2, 5, 8, 11], [0, 1, 4, 7, 10], [0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21], [9, 13, 16, 18, 20]]
+t2m_left_hand_chain = [[20, 22, 23, 24], [20, 34, 35, 36], [20, 25, 26, 27], [20, 31, 32, 33], [20, 28, 29, 30]]
+t2m_right_hand_chain = [[21, 43, 44, 45], [21, 46, 47, 48], [21, 40, 41, 42], [21, 37, 38, 39], [21, 49, 50, 51]]
+
+kit_raw_offsets = np.array(
+    [[0, 0, 0], [0, 1, 0], [0, 1, 0], [0, 1, 0], [0, 1, 0],
+     [1, 0, 0], [0, -1, 0], [0, -1, 0], [-1, 0, 0], [0, -1, 0],
+     [0, -1, 0], [1, 0, 0], [0, -1, 0], [0, -1, 0], [0, 0, 1],
+     [0, 0, 1], [-1, 0, 0], [0, -1, 0], [0, -1, 0], [0, 0, 1],
+     [0, 0, 1]])
+t2m_raw_offsets = np.array([[0,0,0], [1,0,0], [-1,0,0], [0,1,0], [0,-1,0],
+                            [0,-1,0], [0,1,0], [0,-1,0], [0,-1,0], [0,1,0],
+                            [0,0,1], [0,0,1], [0,1,0], [1,0,0], [-1,0,0],
+                            [0,0,1], [0,-1,0], [0,-1,0], [0,-1,0], [0,-1,0],
+                            [0,-1,0], [0,-1,0]])
+
+#################################################################################
+#                                  Joints Revert                                #
+#################################################################################
+def qinv(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    mask = torch.ones_like(q)
+    mask[..., 1:] = -mask[..., 1:]
+    return q * mask
+
+
+def qrot(q, v):
+    """
+    Rotate vector(s) v about the rotation described by quaternion(s) q.
+    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
+    where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 3).
+    """
+    assert q.shape[-1] == 4
+    assert v.shape[-1] == 3
+    assert q.shape[:-1] == v.shape[:-1]
+
+    original_shape = list(v.shape)
+    # print(q.shape)
+    q = q.contiguous().view(-1, 4)
+    v = v.contiguous().view(-1, 3)
+
+    qvec = q[:, 1:]
+    uv = torch.cross(qvec, v, dim=1)
+    uuv = torch.cross(qvec, uv, dim=1)
+    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)
+
+
+def recover_root_rot_pos(data):
+    rot_vel = data[..., 0]
+    r_rot_ang = torch.zeros_like(rot_vel).to(data.device)
+    '''Get Y-axis rotation from rotation velocity'''
+    r_rot_ang[..., 1:] = rot_vel[..., :-1]
+    r_rot_ang = torch.cumsum(r_rot_ang, dim=-1)
+
+    r_rot_quat = torch.zeros(data.shape[:-1] + (4,)).to(data.device)
+    r_rot_quat[..., 0] = torch.cos(r_rot_ang)
+    r_rot_quat[..., 2] = torch.sin(r_rot_ang)
+
+    r_pos = torch.zeros(data.shape[:-1] + (3,)).to(data.device)
+    r_pos[..., 1:, [0, 2]] = data[..., :-1, 1:3]
+    '''Add Y-axis rotation to root position'''
+    r_pos = qrot(qinv(r_rot_quat), r_pos)
+
+    r_pos = torch.cumsum(r_pos, dim=-2)
+
+    r_pos[..., 1] = data[..., 3]
+    return r_rot_quat, r_pos
+
+
+def recover_from_ric(data, joints_num):
+    r_rot_quat, r_pos = recover_root_rot_pos(data)
+    positions = data[..., 4:(joints_num - 1) * 3 + 4]
+    positions = positions.view(positions.shape[:-1] + (-1, 3))
+
+    '''Add Y-axis rotation to local joints'''
+    positions = qrot(qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4,)), positions)
+
+    '''Add root XZ to joints'''
+    positions[..., 0] += r_pos[..., 0:1]
+    positions[..., 2] += r_pos[..., 2:3]
+
+    '''Concate root and joints'''
+    positions = torch.cat([r_pos.unsqueeze(-2), positions], dim=-2)
+
+    return positions
+
+#################################################################################
+#                                 Motion Plotting                               #
+#################################################################################
+def plot_3d_motion(save_path, kinematic_tree, joints, title, figsize=(10, 10), fps=120, radius=4, save_frames_dir=None):
+    # Ensure Agg backend is used (already set at module level, but ensure it's active)
+    if matplotlib.get_backend() != 'Agg':
+        matplotlib.use('Agg')
+
+    title_sp = title.split(' ')
+    if len(title_sp) > 20:
+        title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:20]), ' '.join(title_sp[20:])])
+    elif len(title_sp) > 10:
+        title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:])])
+
+    def init():
+        ax.set_xlim3d([-radius / 2, radius / 2])
+        ax.set_ylim3d([0, radius])
+        ax.set_zlim3d([0, radius])
+        fig.suptitle(title, fontsize=20)
+        ax.grid(b=False)
+
+    def plot_xzPlane(minx, maxx, miny, minz, maxz):
+        verts = [
+            [minx, miny, minz],
+            [minx, miny, maxz],
+            [maxx, miny, maxz],
+            [maxx, miny, minz]
+        ]
+        xz_plane = Poly3DCollection([verts])
+        xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))
+        ax.add_collection3d(xz_plane)
+
+    # Ensure joints is in the correct format: (seq_len, num_joints, 3)
+    joints = np.array(joints)
+    if joints.ndim == 3:
+        # Already in (seq_len, num_joints, 3) format
+        data = joints.copy()
+    elif joints.ndim == 2:
+        # If 2D, reshape to (seq_len, num_joints, 3)
+        # Assume it's (seq_len * num_joints, 3) or (seq_len, num_joints * 3)
+        if joints.shape[1] == 3:
+            # (seq_len * num_joints, 3) - need to infer num_joints
+            # For t2m, we expect 22 joints
+            num_joints = 22
+            seq_len = joints.shape[0] // num_joints
+            data = joints.reshape(seq_len, num_joints, 3)
+        else:
+            # (seq_len, num_joints * 3)
+            num_joints = joints.shape[1] // 3
+            data = joints.reshape(len(joints), num_joints, 3)
+    else:
+        raise ValueError(f"Invalid joints shape: {joints.shape}, expected (seq_len, num_joints, 3)")
+    
+    # Check if data is valid
+    if data.size == 0:
+        raise ValueError("Invalid motion data: data is empty")
+    
+    fig = plt.figure(figsize=figsize)
+    ax = p3.Axes3D(fig)
+    init()
+    MINS = data.min(axis=0).min(axis=0)
+    MAXS = data.max(axis=0).max(axis=0)
+    colors = ['red', 'blue', 'black', 'red', 'blue',
+              'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue',
+              'darkred', 'darkred', 'darkred', 'darkred', 'darkred']
+    frame_number = data.shape[0]
+
+    height_offset = MINS[1]
+    data[:, :, 1] -= height_offset
+    trajec = data[:, 0, [0, 2]]
+
+    data[..., 0] -= data[:, 0:1, 0]
+    data[..., 2] -= data[:, 0:1, 2]
+    
+    # Recompute bounds after centering
+    MINS = data.min(axis=0).min(axis=0)
+    MAXS = data.max(axis=0).max(axis=0)
+    # Add some padding
+    center = (MINS + MAXS) / 2
+    ranges = MAXS - MINS
+    # Ensure we have a minimum range to avoid issues with very small or zero ranges
+    min_range = 0.1  # Minimum range for each axis
+    ranges = np.maximum(ranges, min_range)
+    max_range = max(ranges) * 1.2  # 20% padding
+    plot_mins = center - max_range / 2
+    plot_maxs = center + max_range / 2
+    
+
+    def update(index):
+        # Clear axes properly
+        ax.cla()
+        # Reapply title
+        fig.suptitle(title, fontsize=20)
+        # Reapply view settings and limits based on actual data bounds
+        ax.set_xlim3d([plot_mins[0], plot_maxs[0]])
+        ax.set_ylim3d([plot_mins[1], plot_maxs[1]])
+        ax.set_zlim3d([plot_mins[2], plot_maxs[2]])
+        ax.view_init(elev=120, azim=-90)
+        ax.dist = 7.5
+        ax.grid(False)
+        plot_xzPlane(plot_mins[0] - trajec[index, 0], plot_maxs[0] - trajec[index, 0], 0, 
+                     plot_mins[2] - trajec[index, 1], plot_maxs[2] - trajec[index, 1])
+
+        if index > 1:
+            ax.plot3D(trajec[:index, 0] - trajec[index, 0], np.zeros_like(trajec[:index, 0]),
+                      trajec[:index, 1] - trajec[index, 1], linewidth=1.0,
+                      color='blue')
+
+        for i, (chain, color) in enumerate(zip(kinematic_tree, colors)):
+            if i < len(colors):
+                if i < 5:
+                    linewidth = 4.0
+                else:
+                    linewidth = 2.0
+                # Ensure chain indices are valid
+                valid_chain = [idx for idx in chain if idx < data.shape[1]]
+                if len(valid_chain) > 1:
+                    ax.plot3D(data[index, valid_chain, 0], data[index, valid_chain, 1], data[index, valid_chain, 2], 
+                             linewidth=linewidth, color=color)
+
+        plt.axis('off')
+        ax.set_xticklabels([])
+        ax.set_yticklabels([])
+        ax.set_zticklabels([])
+    
+    def render_frame(frame_idx):
+        """Render a single frame and return as PIL Image"""
+        # Create a fresh figure for each frame to avoid 3D projection issues
+        frame_fig = plt.figure(figsize=figsize, facecolor='white')
+        frame_ax = frame_fig.add_subplot(111, projection='3d')
+        frame_fig.suptitle(title, fontsize=20)
+        
+        # Set limits and view
+        frame_ax.set_xlim3d([plot_mins[0], plot_maxs[0]])
+        frame_ax.set_ylim3d([plot_mins[1], plot_maxs[1]])
+        frame_ax.set_zlim3d([plot_mins[2], plot_maxs[2]])
+        frame_ax.view_init(elev=120, azim=-90)
+        frame_ax.dist = 7.5
+        frame_ax.grid(False)
+        
+        # Plot ground plane
+        minx = plot_mins[0] - trajec[frame_idx, 0]
+        maxx = plot_maxs[0] - trajec[frame_idx, 0]
+        minz = plot_mins[2] - trajec[frame_idx, 1]
+        maxz = plot_maxs[2] - trajec[frame_idx, 1]
+        verts = [[minx, 0, minz], [minx, 0, maxz], [maxx, 0, maxz], [maxx, 0, minz]]
+        xz_plane = Poly3DCollection([verts])
+        xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))
+        frame_ax.add_collection3d(xz_plane)
+        
+        # Plot trajectory
+        if frame_idx > 1:
+            frame_ax.plot3D(trajec[:frame_idx, 0] - trajec[frame_idx, 0], np.zeros_like(trajec[:frame_idx, 0]),
+                      trajec[:frame_idx, 1] - trajec[frame_idx, 1], linewidth=1.0,
+                      color='blue')
+        
+        # Plot skeleton
+        for i, (chain, color) in enumerate(zip(kinematic_tree, colors)):
+            if i < len(colors):
+                if i < 5:
+                    linewidth = 4.0
+                else:
+                    linewidth = 2.0
+                valid_chain = [idx for idx in chain if idx < data.shape[1]]
+                if len(valid_chain) > 1:
+                    frame_ax.plot3D(data[frame_idx, valid_chain, 0], data[frame_idx, valid_chain, 1], 
+                                   data[frame_idx, valid_chain, 2], linewidth=linewidth, color=color)
+        
+        plt.axis('off')
+        frame_ax.set_xticklabels([])
+        frame_ax.set_yticklabels([])
+        frame_ax.set_zticklabels([])
+        
+        # Convert to image - copy data so it's independent of the buffer
+        buf = io.BytesIO()
+        frame_fig.savefig(buf, format='png', dpi=100, bbox_inches='tight', facecolor='white', edgecolor='none')
+        buf.seek(0)
+        # Create image and convert to RGB for GIF compatibility
+        img = Image.open(buf)
+        if img.mode != 'RGB':
+            img = img.convert('RGB')
+        img_copy = img.copy()  # Create a copy that's independent of the buffer
+        buf.close()
+        plt.close(frame_fig)
+        return img_copy
+
+    # Always use frame-by-frame rendering for reliability (works with both GIF and MP4)
+    # This ensures 3D plots render correctly
+    actual_path = save_path
+    
+    # Note: We'll use frame-by-frame rendering for both MP4 and GIF
+    # imageio-ffmpeg will be used for MP4 if available
+    
+    # Use frame-by-frame rendering (works reliably for GIF)
+    if Image is None:
+        raise RuntimeError("PIL/Pillow is required for GIF generation. Please install: pip install Pillow")
+    
+    # Use provided frames directory or create one for debugging
+    frames_dir = save_frames_dir
+    if frames_dir is not None:
+        os.makedirs(frames_dir, exist_ok=True)
+        print(f"Saving individual frames to: {frames_dir}")
+    
+    frames = []
+    print(f"Rendering {frame_number} frames...")
+    for i in range(frame_number):
+        if (i + 1) % 20 == 0:
+            print(f"  Frame {i+1}/{frame_number}")
+        frame_img = render_frame(i)
+        frames.append(frame_img)
+        
+        # Save individual frame as PNG for debugging
+        if frames_dir is not None:
+            frame_path = os.path.join(frames_dir, f"frame_{i:04d}.png")
+            frame_img.save(frame_path)
+            if i == 0 or i == frame_number - 1:
+                print(f"    Saved frame {i} to {frame_path}")
+    
+    # Save video - prefer MP4 if imageio-ffmpeg is available, otherwise GIF
+    if len(frames) > 0:
+        # Ensure all frames are in the same mode and size
+        frames_rgb = []
+        first_frame = frames[0]
+        if first_frame.mode != 'RGB':
+            first_frame = first_frame.convert('RGB')
+        
+        # Get size from first frame
+        target_size = first_frame.size
+        frames_rgb.append(first_frame)
+        
+        # Convert and resize all other frames to match
+        for frame in frames[1:]:
+            if frame.mode != 'RGB':
+                frame = frame.convert('RGB')
+            # Ensure all frames are the same size
+            if frame.size != target_size:
+                frame = frame.resize(target_size, Image.Resampling.LANCZOS)
+            frames_rgb.append(frame)
+        
+        # Convert PIL Images to numpy arrays for imageio
+        frame_arrays = [np.array(frame) for frame in frames_rgb]
+        
+        # Try to save as MP4 first (better quality and compatibility)
+        if actual_path.endswith('.mp4'):
+            if imageio is not None:
+                try:
+                    # Use imageio-ffmpeg for MP4 (automatically uses ffmpeg if available)
+                    imageio.mimsave(actual_path, frame_arrays, fps=fps, codec='libx264', quality=8)
+                    print(f"Saved {len(frames_rgb)} frames to MP4 using imageio-ffmpeg: {actual_path}")
+                except Exception as e:
+                    print(f"Error saving MP4 with imageio-ffmpeg: {e}")
+                    print("Falling back to GIF format...")
+                    # Fall back to GIF
+                    base_path = os.path.splitext(actual_path)[0]
+                    actual_path = base_path + '.gif'
+                    if imageio is not None:
+                        try:
+                            imageio.mimsave(actual_path, frame_arrays, duration=1.0/fps, loop=0)
+                            print(f"Saved {len(frames_rgb)} frames to GIF using imageio: {actual_path}")
+                        except Exception as e2:
+                            print(f"imageio GIF failed, using PIL: {e2}")
+                            frames_rgb[0].save(
+                                actual_path,
+                                save_all=True,
+                                append_images=frames_rgb[1:],
+                                duration=int(1000 / fps),
+                                loop=0,
+                                optimize=False
+                            )
+                    else:
+                        frames_rgb[0].save(
+                            actual_path,
+                            save_all=True,
+                            append_images=frames_rgb[1:],
+                            duration=int(1000 / fps),
+                            loop=0,
+                            optimize=False
+                        )
+            else:
+                print("imageio not available, cannot save MP4. Falling back to GIF...")
+                base_path = os.path.splitext(actual_path)[0]
+                actual_path = base_path + '.gif'
+                frames_rgb[0].save(
+                    actual_path,
+                    save_all=True,
+                    append_images=frames_rgb[1:],
+                    duration=int(1000 / fps),
+                    loop=0,
+                    optimize=False
+                )
+        elif actual_path.endswith('.gif'):
+            # Save as GIF
+            if imageio is not None:
+                try:
+                    imageio.mimsave(actual_path, frame_arrays, duration=1.0/fps, loop=0)
+                    print(f"Saved {len(frames_rgb)} frames to GIF using imageio: {actual_path}")
+                except Exception as e:
+                    print(f"imageio failed, using PIL: {e}")
+                    frames_rgb[0].save(
+                        actual_path,
+                        save_all=True,
+                        append_images=frames_rgb[1:],
+                        duration=int(1000 / fps),
+                        loop=0,
+                        optimize=False
+                    )
+            else:
+                frames_rgb[0].save(
+                    actual_path,
+                    save_all=True,
+                    append_images=frames_rgb[1:],
+                    duration=int(1000 / fps),
+                    loop=0,
+                    optimize=False
+                )
+    
+    if frames_dir is not None:
+        print(f"Saved {len(frames)} individual frames to {frames_dir}")
+    
+    return actual_path

utils/quaternion.py

@@ -0,0 +1,519 @@
+# Copyright (c) 2018-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+import torch
+import numpy as np
+
+_EPS4 = np.finfo(float).eps * 4.0
+
+_FLOAT_EPS = np.finfo(np.float64).eps
+
+# PyTorch-backed implementations
+def qinv(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    mask = torch.ones_like(q)
+    mask[..., 1:] = -mask[..., 1:]
+    return q * mask
+
+
+def qinv_np(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    return qinv(torch.from_numpy(q).float()).numpy()
+
+
+def qnormalize(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    return q / torch.norm(q, dim=-1, keepdim=True)
+
+
+def qmul(q, r):
+    """
+    Multiply quaternion(s) q with quaternion(s) r.
+    Expects two equally-sized tensors of shape (*, 4), where * denotes any number of dimensions.
+    Returns q*r as a tensor of shape (*, 4).
+    """
+    assert q.shape[-1] == 4
+    assert r.shape[-1] == 4
+
+    original_shape = q.shape
+
+    # Compute outer product
+    terms = torch.bmm(r.view(-1, 4, 1), q.view(-1, 1, 4))
+
+    w = terms[:, 0, 0] - terms[:, 1, 1] - terms[:, 2, 2] - terms[:, 3, 3]
+    x = terms[:, 0, 1] + terms[:, 1, 0] - terms[:, 2, 3] + terms[:, 3, 2]
+    y = terms[:, 0, 2] + terms[:, 1, 3] + terms[:, 2, 0] - terms[:, 3, 1]
+    z = terms[:, 0, 3] - terms[:, 1, 2] + terms[:, 2, 1] + terms[:, 3, 0]
+    return torch.stack((w, x, y, z), dim=1).view(original_shape)
+
+
+def qrot(q, v):
+    """
+    Rotate vector(s) v about the rotation described by quaternion(s) q.
+    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
+    where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 3).
+    """
+    assert q.shape[-1] == 4
+    assert v.shape[-1] == 3
+    assert q.shape[:-1] == v.shape[:-1]
+
+    original_shape = list(v.shape)
+    # print(q.shape)
+    q = q.contiguous().view(-1, 4)
+    v = v.contiguous().view(-1, 3)
+
+    qvec = q[:, 1:]
+    uv = torch.cross(qvec, v, dim=1)
+    uuv = torch.cross(qvec, uv, dim=1)
+    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)
+
+
+def qeuler(q, order, epsilon=0, deg=True, follow_order=True):
+    """
+    Convert quaternion(s) q to Euler angles.
+    Expects a tensor of shape (*, 4), where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 3).
+    """
+    assert q.shape[-1] == 4
+
+    original_shape = list(q.shape)
+    original_shape[-1] = 3
+    q = q.view(-1, 4)
+
+    q0 = q[:, 0]
+    q1 = q[:, 1]
+    q2 = q[:, 2]
+    q3 = q[:, 3]
+
+    if order == 'xyz':
+        x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        y = torch.asin(torch.clamp(2 * (q1 * q3 + q0 * q2), -1 + epsilon, 1 - epsilon))
+        z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
+    elif order == 'yzx':
+        x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+        y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3))
+        z = torch.asin(torch.clamp(2 * (q1 * q2 + q0 * q3), -1 + epsilon, 1 - epsilon))
+    elif order == 'zxy':
+        x = torch.asin(torch.clamp(2 * (q0 * q1 + q2 * q3), -1 + epsilon, 1 - epsilon))
+        y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q1 * q1 + q3 * q3))
+    elif order == 'xzy':
+        x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+        y = torch.atan2(2 * (q0 * q2 + q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3))
+        z = torch.asin(torch.clamp(2 * (q0 * q3 - q1 * q2), -1 + epsilon, 1 - epsilon))
+    elif order == 'yxz':
+        x = torch.asin(torch.clamp(2 * (q0 * q1 - q2 * q3), -1 + epsilon, 1 - epsilon))
+        y = torch.atan2(2 * (q1 * q3 + q0 * q2), 1 - 2 * (q1 * q1 + q2 * q2))
+        z = torch.atan2(2 * (q1 * q2 + q0 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+    elif order == 'zyx':
+        x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        y = torch.asin(torch.clamp(2 * (q0 * q2 - q1 * q3), -1 + epsilon, 1 - epsilon))
+        z = torch.atan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
+    else:
+        raise
+    resdict = {"x":x, "y":y, "z":z}
+
+    # print(order)
+    reslist = [resdict[order[i]] for i in range(len(order))] if follow_order else [x, y, z]
+    # print(reslist)
+    if deg:
+        return torch.stack(reslist, dim=1).view(original_shape) * 180 / np.pi
+    else:
+        return torch.stack(reslist, dim=1).view(original_shape)
+
+
+# Numpy-backed implementations
+
+def qmul_np(q, r):
+    q = torch.from_numpy(q).contiguous().float()
+    r = torch.from_numpy(r).contiguous().float()
+    return qmul(q, r).numpy()
+
+
+def qrot_np(q, v):
+    q = torch.from_numpy(q).contiguous().float()
+    v = torch.from_numpy(v).contiguous().float()
+    return qrot(q, v).numpy()
+
+
+def qeuler_np(q, order, epsilon=0, use_gpu=False):
+    if use_gpu:
+        q = torch.from_numpy(q).cuda().float()
+        return qeuler(q, order, epsilon).cpu().numpy()
+    else:
+        q = torch.from_numpy(q).contiguous().float()
+        return qeuler(q, order, epsilon).numpy()
+
+
+def qfix(q):
+    """
+    Enforce quaternion continuity across the time dimension by selecting
+    the representation (q or -q) with minimal distance (or, equivalently, maximal dot product)
+    between two consecutive frames.
+
+    Expects a tensor of shape (L, J, 4), where L is the sequence length and J is the number of joints.
+    Returns a tensor of the same shape.
+    """
+    assert len(q.shape) == 3
+    assert q.shape[-1] == 4
+
+    result = q.copy()
+    dot_products = np.sum(q[1:] * q[:-1], axis=2)
+    mask = dot_products < 0
+    mask = (np.cumsum(mask, axis=0) % 2).astype(bool)
+    result[1:][mask] *= -1
+    return result
+
+
+def euler2quat(e, order, deg=True):
+    """
+    Convert Euler angles to quaternions.
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+
+    e = e.view(-1, 3)
+
+    ## if euler angles in degrees
+    if deg:
+        e = e * np.pi / 180.
+
+    x = e[:, 0]
+    y = e[:, 1]
+    z = e[:, 2]
+
+    rx = torch.stack((torch.cos(x / 2), torch.sin(x / 2), torch.zeros_like(x), torch.zeros_like(x)), dim=1)
+    ry = torch.stack((torch.cos(y / 2), torch.zeros_like(y), torch.sin(y / 2), torch.zeros_like(y)), dim=1)
+    rz = torch.stack((torch.cos(z / 2), torch.zeros_like(z), torch.zeros_like(z), torch.sin(z / 2)), dim=1)
+
+    result = None
+    for coord in order:
+        if coord == 'x':
+            r = rx
+        elif coord == 'y':
+            r = ry
+        elif coord == 'z':
+            r = rz
+        else:
+            raise
+        if result is None:
+            result = r
+        else:
+            result = qmul(result, r)
+
+    # Reverse antipodal representation to have a non-negative "w"
+    if order in ['xyz', 'yzx', 'zxy']:
+        result *= -1
+
+    return result.view(original_shape)
+
+
+def expmap_to_quaternion(e):
+    """
+    Convert axis-angle rotations (aka exponential maps) to quaternions.
+    Stable formula from "Practical Parameterization of Rotations Using the Exponential Map".
+    Expects a tensor of shape (*, 3), where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 4).
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+    e = e.reshape(-1, 3)
+
+    theta = np.linalg.norm(e, axis=1).reshape(-1, 1)
+    w = np.cos(0.5 * theta).reshape(-1, 1)
+    xyz = 0.5 * np.sinc(0.5 * theta / np.pi) * e
+    return np.concatenate((w, xyz), axis=1).reshape(original_shape)
+
+
+def euler_to_quaternion(e, order):
+    """
+    Convert Euler angles to quaternions.
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+
+    e = e.reshape(-1, 3)
+
+    x = e[:, 0]
+    y = e[:, 1]
+    z = e[:, 2]
+
+    rx = np.stack((np.cos(x / 2), np.sin(x / 2), np.zeros_like(x), np.zeros_like(x)), axis=1)
+    ry = np.stack((np.cos(y / 2), np.zeros_like(y), np.sin(y / 2), np.zeros_like(y)), axis=1)
+    rz = np.stack((np.cos(z / 2), np.zeros_like(z), np.zeros_like(z), np.sin(z / 2)), axis=1)
+
+    result = None
+    for coord in order:
+        if coord == 'x':
+            r = rx
+        elif coord == 'y':
+            r = ry
+        elif coord == 'z':
+            r = rz
+        else:
+            raise
+        if result is None:
+            result = r
+        else:
+            result = qmul_np(result, r)
+
+    # Reverse antipodal representation to have a non-negative "w"
+    if order in ['xyz', 'yzx', 'zxy']:
+        result *= -1
+
+    return result.reshape(original_shape)
+
+
+def quaternion_to_matrix(quaternions):
+    """
+    Convert rotations given as quaternions to rotation matrices.
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def quaternion_to_matrix_np(quaternions):
+    q = torch.from_numpy(quaternions).contiguous().float()
+    return quaternion_to_matrix(q).numpy()
+
+
+def quaternion_to_cont6d_np(quaternions):
+    rotation_mat = quaternion_to_matrix_np(quaternions)
+    cont_6d = np.concatenate([rotation_mat[..., 0], rotation_mat[..., 1]], axis=-1)
+    return cont_6d
+
+
+def quaternion_to_cont6d(quaternions):
+    rotation_mat = quaternion_to_matrix(quaternions)
+    cont_6d = torch.cat([rotation_mat[..., 0], rotation_mat[..., 1]], dim=-1)
+    return cont_6d
+
+
+def cont6d_to_matrix(cont6d):
+    assert cont6d.shape[-1] == 6, "The last dimension must be 6"
+    x_raw = cont6d[..., 0:3]
+    y_raw = cont6d[..., 3:6]
+
+    x = x_raw / torch.norm(x_raw, dim=-1, keepdim=True)
+    z = torch.cross(x, y_raw, dim=-1)
+    z = z / torch.norm(z, dim=-1, keepdim=True)
+
+    y = torch.cross(z, x, dim=-1)
+
+    x = x[..., None]
+    y = y[..., None]
+    z = z[..., None]
+
+    mat = torch.cat([x, y, z], dim=-1)
+    return mat
+
+
+def cont6d_to_matrix_np(cont6d):
+    q = torch.from_numpy(cont6d).contiguous().float()
+    return cont6d_to_matrix(q).numpy()
+
+
+def qpow(q0, t, dtype=torch.float):
+    ''' q0 : tensor of quaternions
+    t: tensor of powers
+    '''
+    q0 = qnormalize(q0)
+    theta0 = torch.acos(q0[..., 0])
+
+    ## if theta0 is close to zero, add epsilon to avoid NaNs
+    mask = (theta0 <= 10e-10) * (theta0 >= -10e-10)
+    theta0 = (1 - mask) * theta0 + mask * 10e-10
+    v0 = q0[..., 1:] / torch.sin(theta0).view(-1, 1)
+
+    if isinstance(t, torch.Tensor):
+        q = torch.zeros(t.shape + q0.shape)
+        theta = t.view(-1, 1) * theta0.view(1, -1)
+    else:  ## if t is a number
+        q = torch.zeros(q0.shape)
+        theta = t * theta0
+
+    q[..., 0] = torch.cos(theta)
+    q[..., 1:] = v0 * torch.sin(theta).unsqueeze(-1)
+
+    return q.to(dtype)
+
+
+def qslerp(q0, q1, t):
+    '''
+    q0: starting quaternion
+    q1: ending quaternion
+    t: array of points along the way
+
+    Returns:
+    Tensor of Slerps: t.shape + q0.shape
+    '''
+
+    q0 = qnormalize(q0)
+    q1 = qnormalize(q1)
+    q_ = qpow(qmul(q1, qinv(q0)), t)
+
+    return qmul(q_,
+                q0.contiguous().view(torch.Size([1] * len(t.shape)) + q0.shape).expand(t.shape + q0.shape).contiguous())
+
+
+def qbetween(v0, v1):
+    '''
+    find the quaternion used to rotate v0 to v1
+    '''
+    assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)'
+    assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)'
+
+    v = torch.cross(v0, v1)
+    w = torch.sqrt((v0 ** 2).sum(dim=-1, keepdim=True) * (v1 ** 2).sum(dim=-1, keepdim=True)) + (v0 * v1).sum(dim=-1,
+                                                                                                              keepdim=True)
+    return qnormalize(torch.cat([w, v], dim=-1))
+
+
+def qbetween_np(v0, v1):
+    '''
+    find the quaternion used to rotate v0 to v1
+    '''
+    assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)'
+    assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)'
+
+    v0 = torch.from_numpy(v0).float()
+    v1 = torch.from_numpy(v1).float()
+    return qbetween(v0, v1).numpy()
+
+
+def lerp(p0, p1, t):
+    if not isinstance(t, torch.Tensor):
+        t = torch.Tensor([t])
+
+    new_shape = t.shape + p0.shape
+    new_view_t = t.shape + torch.Size([1] * len(p0.shape))
+    new_view_p = torch.Size([1] * len(t.shape)) + p0.shape
+    p0 = p0.view(new_view_p).expand(new_shape)
+    p1 = p1.view(new_view_p).expand(new_shape)
+    t = t.view(new_view_t).expand(new_shape)
+
+    return p0 + t * (p1 - p0)
+
+def matrix_to_quat(R) -> torch.Tensor:
+    '''
+    https://github.com/duolu/pyrotation/blob/master/pyrotation/pyrotation.py
+    Convert a rotation matrix to a unit quaternion.
+    This uses the Shepperd’s method for numerical stability.
+    '''
+
+    # The rotation matrix must be orthonormal
+
+    w2 = (1 + R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2])
+    x2 = (1 + R[..., 0, 0] - R[..., 1, 1] - R[..., 2, 2])
+    y2 = (1 - R[..., 0, 0] + R[..., 1, 1] - R[..., 2, 2])
+    z2 = (1 - R[..., 0, 0] - R[..., 1, 1] + R[..., 2, 2])
+
+    yz = (R[..., 1, 2] + R[..., 2, 1])
+    xz = (R[..., 2, 0] + R[..., 0, 2])
+    xy = (R[..., 0, 1] + R[..., 1, 0])
+
+    wx = (R[..., 2, 1] - R[..., 1, 2])
+    wy = (R[..., 0, 2] - R[..., 2, 0])
+    wz = (R[..., 1, 0] - R[..., 0, 1])
+
+    w = torch.empty_like(x2)
+    x = torch.empty_like(x2)
+    y = torch.empty_like(x2)
+    z = torch.empty_like(x2)
+
+    flagA = (R[..., 2, 2] < 0) * (R[..., 0, 0] > R[..., 1, 1])
+    flagB = (R[..., 2, 2] < 0) * (R[..., 0, 0] <= R[..., 1, 1])
+    flagC = (R[..., 2, 2] >= 0) * (R[..., 0, 0] < -R[..., 1, 1])
+    flagD = (R[..., 2, 2] >= 0) * (R[..., 0, 0] >= -R[..., 1, 1])
+
+    x[flagA] = torch.sqrt(x2[flagA])
+    w[flagA] = wx[flagA] / x[flagA]
+    y[flagA] = xy[flagA] / x[flagA]
+    z[flagA] = xz[flagA] / x[flagA]
+
+    y[flagB] = torch.sqrt(y2[flagB])
+    w[flagB] = wy[flagB] / y[flagB]
+    x[flagB] = xy[flagB] / y[flagB]
+    z[flagB] = yz[flagB] / y[flagB]
+
+    z[flagC] = torch.sqrt(z2[flagC])
+    w[flagC] = wz[flagC] / z[flagC]
+    x[flagC] = xz[flagC] / z[flagC]
+    y[flagC] = yz[flagC] / z[flagC]
+
+    w[flagD] = torch.sqrt(w2[flagD])
+    x[flagD] = wx[flagD] / w[flagD]
+    y[flagD] = wy[flagD] / w[flagD]
+    z[flagD] = wz[flagD] / w[flagD]
+
+    # if R[..., 2, 2] < 0:
+    #
+    #     if R[..., 0, 0] > R[..., 1, 1]:
+    #
+    #         x = torch.sqrt(x2)
+    #         w = wx / x
+    #         y = xy / x
+    #         z = xz / x
+    #
+    #     else:
+    #
+    #         y = torch.sqrt(y2)
+    #         w = wy / y
+    #         x = xy / y
+    #         z = yz / y
+    #
+    # else:
+    #
+    #     if R[..., 0, 0] < -R[..., 1, 1]:
+    #
+    #         z = torch.sqrt(z2)
+    #         w = wz / z
+    #         x = xz / z
+    #         y = yz / z
+    #
+    #     else:
+    #
+    #         w = torch.sqrt(w2)
+    #         x = wx / w
+    #         y = wy / w
+    #         z = wz / w
+
+    res = [w, x, y, z]
+    res = [z.unsqueeze(-1) for z in res]
+
+    return torch.cat(res, dim=-1) / 2
+
+def cont6d_to_quat(cont6d):
+    return matrix_to_quat(cont6d_to_matrix(cont6d))

utils/skeleton.py

@@ -0,0 +1,194 @@
+from utils.quaternion import *
+import scipy.ndimage.filters as filters
+
+class Skeleton(object):
+    def __init__(self, offset, kinematic_tree, device):
+        self.device = device
+        self._raw_offset_np = offset.numpy()
+        self._raw_offset = offset.clone().detach().to(device).float()
+        self._kinematic_tree = kinematic_tree
+        self._offset = None
+        self._parents = [0] * len(self._raw_offset)
+        self._parents[0] = -1
+        for chain in self._kinematic_tree:
+            for j in range(1, len(chain)):
+                self._parents[chain[j]] = chain[j-1]
+
+    def njoints(self):
+        return len(self._raw_offset)
+
+    def offset(self):
+        return self._offset
+
+    def set_offset(self, offsets):
+        self._offset = offsets.clone().detach().to(self.device).float()
+
+    def kinematic_tree(self):
+        return self._kinematic_tree
+
+    def parents(self):
+        return self._parents
+
+    # joints (batch_size, joints_num, 3)
+    def get_offsets_joints_batch(self, joints):
+        assert len(joints.shape) == 3
+        _offsets = self._raw_offset.expand(joints.shape[0], -1, -1).clone()
+        for i in range(1, self._raw_offset.shape[0]):
+            _offsets[:, i] = torch.norm(joints[:, i] - joints[:, self._parents[i]], p=2, dim=1)[:, None] * _offsets[:, i]
+
+        self._offset = _offsets.detach()
+        return _offsets
+
+    # joints (joints_num, 3)
+    def get_offsets_joints(self, joints):
+        assert len(joints.shape) == 2
+        _offsets = self._raw_offset.clone()
+        for i in range(1, self._raw_offset.shape[0]):
+            # print(joints.shape)
+            _offsets[i] = torch.norm(joints[i] - joints[self._parents[i]], p=2, dim=0) * _offsets[i]
+
+        self._offset = _offsets.detach()
+        return _offsets
+
+    # face_joint_idx should follow the order of right hip, left hip, right shoulder, left shoulder
+    # joints (batch_size, joints_num, 3)
+    def inverse_kinematics_np(self, joints, face_joint_idx, smooth_forward=False):
+        assert len(face_joint_idx) == 4
+        '''Get Forward Direction'''
+        l_hip, r_hip, sdr_r, sdr_l = face_joint_idx
+        across1 = joints[:, r_hip] - joints[:, l_hip]
+        across2 = joints[:, sdr_r] - joints[:, sdr_l]
+        across = across1 + across2
+        across = across / (np.sqrt((across**2).sum(axis=-1, keepdims=True)) + 1e-8)
+        # print(across1.shape, across2.shape)
+
+        # forward (batch_size, 3)
+        forward = np.cross(np.array([[0, 1, 0]]), across, axis=-1)
+        if smooth_forward:
+            forward = filters.gaussian_filter1d(forward, 20, axis=0, mode='nearest')
+            # forward (batch_size, 3)
+        forward = forward / (np.sqrt((forward**2).sum(axis=-1))[..., np.newaxis] + 1e-8)
+
+        '''Get Root Rotation'''
+        target = np.array([[0,0,1]]).repeat(len(forward), axis=0)
+        root_quat = qbetween_np(forward, target)
+
+        '''Inverse Kinematics'''
+        # quat_params (batch_size, joints_num, 4)
+        # print(joints.shape[:-1])
+        quat_params = np.zeros(joints.shape[:-1] + (4,))
+        # print(quat_params.shape)
+        root_quat[0] = np.array([[1.0, 0.0, 0.0, 0.0]])
+        quat_params[:, 0] = root_quat
+        # quat_params[0, 0] = np.array([[1.0, 0.0, 0.0, 0.0]])
+        for chain in self._kinematic_tree:
+            R = root_quat
+            for j in range(len(chain) - 1):
+                # (batch, 3)
+                u = self._raw_offset_np[chain[j+1]][np.newaxis,...].repeat(len(joints), axis=0)
+                # print(u.shape)
+                # (batch, 3)
+                v = joints[:, chain[j+1]] - joints[:, chain[j]]
+                v = v / (np.sqrt((v**2).sum(axis=-1, keepdims=True)) + 1e-8)
+                # print(u.shape, v.shape)
+                rot_u_v = qbetween_np(u, v)
+
+                R_loc = qmul_np(qinv_np(R), rot_u_v)
+
+                quat_params[:,chain[j + 1], :] = R_loc
+                R = qmul_np(R, R_loc)
+
+        return quat_params
+
+    # Be sure root joint is at the beginning of kinematic chains
+    def forward_kinematics(self, quat_params, root_pos, skel_joints=None, do_root_R=True):
+        # quat_params (batch_size, joints_num, 4)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(quat_params.shape[0], -1, -1)
+        joints = torch.zeros(quat_params.shape[:-1] + (3,)).to(self.device)
+        joints[:, 0] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                R = quat_params[:, 0]
+            else:
+                R = torch.tensor([[1.0, 0.0, 0.0, 0.0]]).expand(len(quat_params), -1).detach().to(self.device)
+            for i in range(1, len(chain)):
+                R = qmul(R, quat_params[:, chain[i]])
+                offset_vec = offsets[:, chain[i]]
+                joints[:, chain[i]] = qrot(R, offset_vec) + joints[:, chain[i-1]]
+        return joints
+
+    # Be sure root joint is at the beginning of kinematic chains
+    def forward_kinematics_np(self, quat_params, root_pos, skel_joints=None, do_root_R=True):
+        # quat_params (batch_size, joints_num, 4)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            skel_joints = torch.from_numpy(skel_joints)
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(quat_params.shape[0], -1, -1)
+        offsets = offsets.numpy()
+        joints = np.zeros(quat_params.shape[:-1] + (3,))
+        joints[:, 0] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                R = quat_params[:, 0]
+            else:
+                R = np.array([[1.0, 0.0, 0.0, 0.0]]).repeat(len(quat_params), axis=0)
+            for i in range(1, len(chain)):
+                R = qmul_np(R, quat_params[:, chain[i]])
+                offset_vec = offsets[:, chain[i]]
+                joints[:, chain[i]] = qrot_np(R, offset_vec) + joints[:, chain[i - 1]]
+        return joints
+
+    def forward_kinematics_cont6d_np(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True):
+        # cont6d_params (batch_size, joints_num, 6)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            skel_joints = torch.from_numpy(skel_joints)
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(cont6d_params.shape[0], -1, -1)
+        offsets = offsets.numpy()
+        joints = np.zeros(cont6d_params.shape[:-1] + (3,))
+        joints[:, 0] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                matR = cont6d_to_matrix_np(cont6d_params[:, 0])
+            else:
+                matR = np.eye(3)[np.newaxis, :].repeat(len(cont6d_params), axis=0)
+            for i in range(1, len(chain)):
+                matR = np.matmul(matR, cont6d_to_matrix_np(cont6d_params[:, chain[i]]))
+                offset_vec = offsets[:, chain[i]][..., np.newaxis]
+                # print(matR.shape, offset_vec.shape)
+                joints[:, chain[i]] = np.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]]
+        return joints
+
+    def forward_kinematics_cont6d(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True):
+        # cont6d_params (batch_size, joints_num, 6)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            # skel_joints = torch.from_numpy(skel_joints)
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(cont6d_params.shape[0], -1, -1)
+        joints = torch.zeros(cont6d_params.shape[:-1] + (3,)).to(cont6d_params.device)
+        joints[..., 0, :] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                matR = cont6d_to_matrix(cont6d_params[:, 0])
+            else:
+                matR = torch.eye(3).expand((len(cont6d_params), -1, -1)).detach().to(cont6d_params.device)
+            for i in range(1, len(chain)):
+                matR = torch.matmul(matR, cont6d_to_matrix(cont6d_params[:, chain[i]]))
+                offset_vec = offsets[:, chain[i]].unsqueeze(-1)
+                # print(matR.shape, offset_vec.shape)
+                joints[:, chain[i]] = torch.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]]
+        return joints

utils/train_utils.py

@@ -0,0 +1,105 @@
+import torch
+import math
+import time
+import torch.distributed as dist
+import logging
+#################################################################################
+#                                  DDP Functions                               #
+#################################################################################
+def cleanup():
+    dist.destroy_process_group()
+
+#################################################################################
+#                                  Util Functions                               #
+#################################################################################
+def lengths_to_mask(lengths, max_len):
+    # max_len = max(lengths)
+    mask = torch.arange(max_len, device=lengths.device).expand(len(lengths), max_len) < lengths.unsqueeze(1)
+    return mask #(b, len)
+
+
+def get_mask_subset_prob(mask, prob):
+    subset_mask = torch.bernoulli(mask, p=prob) & mask
+    return subset_mask
+
+
+def uniform(shape, device=None):
+    return torch.zeros(shape, device=device).float().uniform_(0, 1)
+
+
+def cosine_schedule(t):
+    return torch.cos(t * math.pi * 0.5)
+
+
+def update_ema(model, ema_model, ema_decay):
+    with torch.no_grad():
+        for ema_param, model_param in zip(ema_model.parameters(), model.parameters()):
+            ema_param.data.mul_(ema_decay).add_(model_param.data, alpha=(1 - ema_decay))
+
+#################################################################################
+#                                Logging Functions                              #
+#################################################################################
+def def_value():
+    return 0.0
+
+
+def create_logger(logging_dir):
+    if dist.get_rank() == 0:  # real logger
+        logging.basicConfig(
+            level=logging.INFO,
+            format='[\033[34m%(asctime)s\033[0m] %(message)s',
+            datefmt='%Y-%m-%d %H:%M:%S',
+            handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
+        )
+        logger = logging.getLogger(__name__)
+    else:  # dummy logger (does nothing)
+        logger = logging.getLogger(__name__)
+        logger.addHandler(logging.NullHandler())
+    return logger
+
+
+def update_lr_warm_up(nb_iter, warm_up_iter, optimizer, lr):
+    current_lr = lr * (nb_iter + 1) / (warm_up_iter + 1)
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = current_lr
+    return current_lr
+
+
+def save(file_name, ep, model, optimizer, scheduler, total_it, name, ema=None):
+    state = {
+        name: model.state_dict(),
+        f"opt_{name}": optimizer.state_dict(),
+        "scheduler": scheduler.state_dict(),
+        'ep': ep,
+        'total_it': total_it,
+    }
+    if ema is not None:
+        mardm_state_dict = model.state_dict()
+        ema_mardm_state_dict = ema.state_dict()
+        clip_weights = [e for e in mardm_state_dict.keys() if e.startswith('clip_model.')]
+        for e in clip_weights:
+            del mardm_state_dict[e]
+            del ema_mardm_state_dict[e]
+        state[name] = mardm_state_dict
+        state[f"ema_{name}"] = ema_mardm_state_dict
+    torch.save(state, file_name)
+
+
+def print_current_loss(start_time, niter_state, total_niters, losses, epoch=None, sub_epoch=None,
+                       inner_iter=None, tf_ratio=None, sl_steps=None):
+    def as_minutes(s):
+        m = math.floor(s / 60)
+        s -= m * 60
+        return '%dm %ds' % (m, s)
+    def time_since(since, percent):
+        now = time.time()
+        s = now - since
+        es = s / percent
+        rs = es - s
+        return '%s (- %s)' % (as_minutes(s), as_minutes(rs))
+    if epoch is not None:
+        print('ep/it:%2d-%4d niter:%6d' % (epoch, inner_iter, niter_state), end=" ")
+    message = ' %s completed:%3d%%)' % (time_since(start_time, niter_state / total_niters), niter_state / total_niters * 100)
+    for k, v in losses.items():
+        message += ' %s: %.4f ' % (k, v)
+    print(message)

utils/wandb_utils.py

@@ -0,0 +1,53 @@
+import wandb
+import torch
+from torchvision.utils import make_grid
+import torch.distributed as dist
+import argparse
+import hashlib
+import math
+
+
+def is_main_process():
+    return dist.get_rank() == 0
+
+def namespace_to_dict(namespace):
+    return {
+        k: namespace_to_dict(v) if isinstance(v, argparse.Namespace) else v
+        for k, v in vars(namespace).items()
+    }
+
+
+def generate_run_id(exp_name):
+    # https://stackoverflow.com/questions/16008670/how-to-hash-a-string-into-8-digits
+    return str(int(hashlib.sha256(exp_name.encode('utf-8')).hexdigest(), 16) % 10 ** 8)
+
+
+def initialize(args, entity, exp_name, project_name):
+    config_dict = namespace_to_dict(args)
+    # wandb.login(key=os.environ["WANDB_KEY"])
+    wandb.init(
+        entity=entity,
+        project=project_name,
+        name=exp_name,
+        config=config_dict,
+        id=generate_run_id(exp_name),
+        resume="allow",
+    )
+
+
+def log(stats, step=None):
+    if is_main_process():
+        wandb.log({k: v for k, v in stats.items()}, step=step)
+
+
+def log_image(sample, step=None):
+    if is_main_process():
+        sample = array2grid(sample)
+        wandb.log({f"samples": wandb.Image(sample), "train_step": step})
+
+
+def array2grid(x):
+    nrow = round(math.sqrt(x.size(0)))
+    x = make_grid(x, nrow=nrow, normalize=True, value_range=(-1,1))
+    x = x.mul(255).add_(0.5).clamp_(0,255).permute(1,2,0).to('cpu', torch.uint8).numpy()
+    return x