import os
import uuid
import subprocess
import shlex
import gradio as gr

MODEL_DIR = "motion-diffusion-model"   # where you copied the repo
OUTPUT_DIR = "outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)

def generate_motion(prompt: str):
    # Unique filename
    run_id = uuid.uuid4().hex[:8]
    out_file = f"{run_id}.mp4"
    out_path = os.path.join(OUTPUT_DIR, out_file)

    # Command-line call to the sample script
    # Adjust flags if needed; this is the typical pattern
    cmd = f"python {MODEL_DIR}/sample/predict.py " \
          f"--prompt \"{prompt}\" " \
          f"--results_dir {OUTPUT_DIR} " \
          f"--num_samples 1"

    # Run the model
    subprocess.run(shlex.split(cmd), check=True)

    # The script will dump something like outputs/sample_0.mp4
    # Rename it to our unique filename
    default_out = os.path.join(OUTPUT_DIR, "sample_0.mp4")
    if os.path.exists(default_out):
        os.replace(default_out, out_path)
    else:
        raise FileNotFoundError(default_out)

    return out_path

demo = gr.Interface(
    fn=generate_motion,
    inputs=gr.Textbox(label="Enter a text prompt"),
    outputs=gr.Video(label="Generated Motion"),
    title="Motion Diffusion Model",
    description="Type some text and watch the 3D human motion!"
)

if __name__ == "__main__":
    demo.launch()