Use the correct sample.generate module for MDM

app.py

@@ -7,10 +7,6 @@ from pathlib import Path
 import traceback
 import subprocess
 
-# Add the motion-diffusion-model repository to the path if it exists
-if Path("motion-diffusion-model").exists():
-    sys.path.append("./motion-diffusion-model")
-
 def ensure_mdm_repo():
     """Ensure the MDM repository is cloned and set up"""
     if not Path("motion-diffusion-model").exists():
@@ -21,7 +17,7 @@ def ensure_mdm_repo():
         print("Setting up the repository...")
         subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
         
-        # Download other necessary files
+        # Add necessary files
         os.chdir("motion-diffusion-model")
         subprocess.run(["bash", "prepare/download_smpl_files.sh"])
         subprocess.run(["bash", "prepare/download_glove.sh"])
@@ -32,61 +28,6 @@ def ensure_mdm_repo():
     if "./motion-diffusion-model" not in sys.path:
         sys.path.append("./motion-diffusion-model")
 
-def run_mdm_generation(text_prompt, motion_length, seed, checkpoint_path):
-    """Run MDM generation using the provided checkpoint"""
-    # Print the current directory and check that checkpoint exists
-    print(f"Current directory: {os.getcwd()}")
-    print(f"Checking checkpoint path: {checkpoint_path}, exists: {os.path.exists(checkpoint_path)}")
-    
-    # Change to the MDM repository directory
-    original_dir = os.getcwd()
-    os.chdir("motion-diffusion-model")
-    
-    # Prepare output directory
-    os.makedirs("../output", exist_ok=True)
-    output_path = f"../output/output_{abs(hash(text_prompt) % 10000)}_{int(motion_length)}_{seed}.mp4"
-    
-    # Create a command to run the MDM sample script with all necessary parameters
-    cmd = [
-        "python", "-m", "sample.sample_text",
-        "--model_path", checkpoint_path,
-        "--text_prompt", text_prompt,
-        "--motion_length", str(motion_length),
-        "--seed", str(int(seed)),
-        "--num_samples", "1",
-        "--dataset", "humanml",  # Specify the dataset
-        "--guidance_scale", "2.5",  # Higher values follow the text more closely
-        "--num_repetitions", "1",
-    ]
-    
-    print(f"Running command: {' '.join(cmd)}")
-    result = subprocess.run(cmd, capture_output=True, text=True)
-    
-    print("Command output:", result.stdout)
-    if result.stderr:
-        print("Command error:", result.stderr)
-    
-    # Check for output files - assume they're saved to samples directory
-    output_files = []
-    if os.path.exists("samples"):
-        for file in os.listdir("samples"):
-            if file.endswith(".mp4"):
-                output_files.append(os.path.join("samples", file))
-                print(f"Found output file: {file}")
-    
-    # Return to the original directory
-    os.chdir(original_dir)
-    
-    # If we found output files, copy the first one to our output path
-    if output_files:
-        os.makedirs(os.path.dirname(output_path), exist_ok=True)
-        subprocess.run(["cp", os.path.join("motion-diffusion-model", output_files[0]), output_path])
-        print(f"Copied output to {output_path}")
-        return output_path
-    
-    print("No output files found.")
-    return None
-
 def text_to_motion(text_prompt, motion_length=3.0, seed=0):
     """Generate motion from text prompt using MDM"""
     try:
@@ -95,19 +36,58 @@ def text_to_motion(text_prompt, motion_length=3.0, seed=0):
         # Ensure the MDM repository is set up
         ensure_mdm_repo()
         
-        # Use the uploaded checkpoint with absolute path
+        # Create output directory
+        os.makedirs("output", exist_ok=True)
+        
+        # Get absolute path to the checkpoint
         checkpoint_path = os.path.abspath("checkpoints/opt000750000.pt")
         print(f"Using checkpoint: {checkpoint_path}")
         
-        # Try to run MDM generation
-        output_path = run_mdm_generation(text_prompt, motion_length, seed, checkpoint_path)
+        # Change to the MDM repository directory
+        original_dir = os.getcwd()
+        os.chdir("motion-diffusion-model")
+        
+        # Create the command to run - based on official examples
+        cmd = [
+            "python", 
+            "-m", "sample.generate",  # The correct entry point
+            "--model_path", checkpoint_path,
+            "--text_prompt", text_prompt,
+            "--motion_length", str(motion_length),
+            "--seed", str(int(seed)),
+            "--num_samples", "1",     # Generate just one sample
+            "--num_repetitions", "1"  # With one repetition
+        ]
+        
+        print(f"Running command: {' '.join(cmd)}")
+        result = subprocess.run(cmd, capture_output=True, text=True)
         
-        if output_path and os.path.exists(output_path):
-            print(f"Successfully generated motion at {output_path}")
+        # Print the output for debugging
+        print("Command output:", result.stdout)
+        if result.stderr:
+            print("Command error:", result.stderr)
+        
+        # Check for output files - MDM saves samples in samples directory
+        output_mp4 = None
+        if os.path.exists("samples"):
+            for file in os.listdir("samples"):
+                if file.endswith(".mp4"):
+                    output_mp4 = os.path.join("samples", file)
+                    print(f"Found output file: {output_mp4}")
+                    break
+        
+        # Return to the original directory
+        os.chdir(original_dir)
+        
+        # If we found an output file, copy it to our output directory
+        if output_mp4:
+            output_path = f"output/output_{abs(hash(text_prompt) % 10000)}_{int(motion_length)}_{seed}.mp4"
+            subprocess.run(["cp", os.path.join("motion-diffusion-model", output_mp4), output_path])
+            print(f"Copied output to {output_path}")
             return output_path
-        else:
-            print("MDM generation failed, no output file produced")
-            return None
+        
+        print("No output files found.")
+        return None
         
     except Exception as e:
         print(f"Error generating motion: {str(e)}")
@@ -124,7 +104,7 @@ demo = gr.Interface(
     ],
     outputs=gr.Video(label="Generated Motion"),
     title="Motion Diffusion Model Demo",
-    description="Generate human motions from text descriptions using the opt000750000.pt checkpoint model."
+    description="Generate human motions from text descriptions using the opt000750000.pt checkpoint model. Try prompts like: 'A person walks forward, then turns left', 'A person jumps up and down', or 'A person dances energetically'."
 )
 
 # Launch the app