Avoid shell scripts and use direct Python commands

app.py

@@ -8,6 +8,7 @@ import traceback
 import subprocess
 import glob
 import requests
+import time
 
 def download_checkpoint():
     """Download the recommended checkpoint if not present"""
@@ -64,44 +65,157 @@ def text_to_motion(text_prompt, motion_length=3.0, seed=0):
         
         # Download the recommended checkpoint
         checkpoint_path = download_checkpoint()
+        absolute_checkpoint_path = os.path.abspath(checkpoint_path)
         
-        # Write a simple run script that properly sets up the environment
-        run_script = """
-#!/bin/bash
-cd motion-diffusion-model
-export PYTHONPATH=$PYTHONPATH:$(pwd)
-python -m sample.generate --model_path ../checkpoints/humanml_trans_enc_512.pt --text_prompt "$1" --motion_length $2 --seed $3 --num_samples 1
-"""
+        # Set up environment
+        original_dir = os.getcwd()
+        os.chdir("motion-diffusion-model")
         
-        with open("run_mdm.sh", "w") as f:
-            f.write(run_script)
+        # Add current directory to Python path
+        sys.path.insert(0, os.getcwd())
         
-        subprocess.run(["chmod", "+x", "run_mdm.sh"])
+        # Create a simple visualization script
+        with open("visualize_motion.py", "w") as f:
+            f.write("""
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
+import sys
+import os
+from mpl_toolkits.mplot3d import Axes3D
+
+def visualize_motion(motion_file, output_path):
+    # Load the motion data
+    motion_data = np.load(motion_file)
+    
+    # Get dimensions
+    frames, joints, dims = motion_data.shape
+    
+    # Create figure
+    fig = plt.figure(figsize=(10, 10))
+    ax = fig.add_subplot(111, projection='3d')
+    
+    # Define connections between joints (simplified)
+    connections = [
+        (0, 1), (1, 2), (2, 3),  # Spine
+        (0, 4), (4, 5), (5, 6),  # Left leg
+        (0, 7), (7, 8), (8, 9),  # Right leg
+        (3, 10), (10, 11), (11, 12),  # Left arm
+        (3, 13), (13, 14), (14, 15)   # Right arm
+    ]
+    
+    def update(frame):
+        ax.clear()
         
-        # Run the script
-        cmd = ["./run_mdm.sh", text_prompt, str(motion_length), str(int(seed))]
-        print(f"Running command: {' '.join(cmd)}")
-        result = subprocess.run(cmd, capture_output=True, text=True)
+        # Set axis limits
+        max_range = np.max(np.abs(motion_data))
+        ax.set_xlim([-max_range, max_range])
+        ax.set_ylim([-max_range, max_range])
+        ax.set_zlim([-max_range, max_range])
         
-        # Print the output for debugging
-        print("Command output:", result.stdout)
-        if result.stderr:
-            print("Command error:", result.stderr)
+        # Plot joints
+        ax.scatter(motion_data[frame, :, 0], 
+                  motion_data[frame, :, 1], 
+                  motion_data[frame, :, 2], c='b', marker='o')
         
-        # Check for output files
-        print("Checking for output files:")
-        for root, dirs, files in os.walk("."):
-            for file in files:
-                if file.endswith(".mp4"):
-                    path = os.path.join(root, file)
-                    print(f"Found video file: {path}")
-                    return path
+        # Plot connections
+        for start, end in connections:
+            if start < joints and end < joints:
+                ax.plot([motion_data[frame, start, 0], motion_data[frame, end, 0]],
+                       [motion_data[frame, start, 1], motion_data[frame, end, 1]],
+                       [motion_data[frame, start, 2], motion_data[frame, end, 2]], 'r-')
         
-        print("No MP4 output files found.")
+        ax.set_title(f"Frame {frame}")
+        return ax
+    
+    # Create animation
+    anim = FuncAnimation(fig, update, frames=min(100, frames), interval=1000/30)
+    
+    # Save animation
+    anim.save(output_path, writer='ffmpeg', fps=30)
+    plt.close()
+    
+    return output_path
+
+if __name__ == "__main__":
+    if len(sys.argv) < 3:
+        print("Usage: python visualize_motion.py <motion_file> <output_path>")
+        sys.exit(1)
+    
+    motion_file = sys.argv[1]
+    output_path = sys.argv[2]
+    
+    visualize_motion(motion_file, output_path)
+""")
+        
+        # Run the generation directly
+        print("Running MDM generation...")
+        generation_cmd = [
+            "python", "-c", 
+            f"""
+import sys
+sys.path.insert(0, '.')
+from sample.generate import generate
+import os
+
+# Run the generation
+motion_data = generate(
+    model_path='{absolute_checkpoint_path}',
+    text_prompt='{text_prompt}',
+    motion_length={motion_length},
+    seeds=[{int(seed)}],
+    num_samples=1
+)
+
+# Save the motion data for visualization
+import numpy as np
+os.makedirs('output', exist_ok=True)
+np.save('output/motion_data.npy', motion_data[0])
+print("Motion data saved to output/motion_data.npy")
+"""
+        ]
+        
+        print(f"Running generation command...")
+        gen_result = subprocess.run(generation_cmd, capture_output=True, text=True)
+        
+        print("Generation output:", gen_result.stdout)
+        if gen_result.stderr:
+            print("Generation error:", gen_result.stderr)
+        
+        # Check if motion data was generated
+        if Path("output/motion_data.npy").exists():
+            print("Motion data generated successfully!")
+            
+            # Visualize the motion
+            print("Visualizing motion...")
+            viz_cmd = ["python", "visualize_motion.py", "output/motion_data.npy", "../output.mp4"]
+            viz_result = subprocess.run(viz_cmd, capture_output=True, text=True)
+            
+            print("Visualization output:", viz_result.stdout)
+            if viz_result.stderr:
+                print("Visualization error:", viz_result.stderr)
+            
+            # Return to original directory
+            os.chdir(original_dir)
+            
+            # Check if the output file exists
+            if Path("output.mp4").exists():
+                print("Animation created successfully!")
+                return "output.mp4"
+        
+        # Return to original directory
+        os.chdir(original_dir)
+        
+        print("No output file generated.")
         return None
     except Exception as e:
         print(f"Error generating motion: {str(e)}")
         print(traceback.format_exc())  # Print the full traceback
+        
+        # Return to original directory if an exception occurred
+        if 'original_dir' in locals():
+            os.chdir(original_dir)
+            
         return None
 
 # Create the Gradio interface