import gradio as gr
import torch
from in2in.utils.plot import plot_3d_motion
from in2in.utils.paramUtil import HML_KINEMATIC_CHAIN
from transformers import AutoModel     

import random
import string

def generate_random_filename(length=10, extension='.mp4'):
    """
    Generates a random file name with the specified length and file extension.
    
    Args:
        length (int): The desired length of the file name (excluding the extension).
        extension (str): The file extension, including the dot (e.g., '.txt', '.jpg', '.pdf').
    
    Returns:
        str: The generated random file name with the specified extension.
    """
    characters = string.ascii_letters + string.digits
    filename = ''.join(random.choice(characters) for _ in range(length))
    return filename + extension

def generate(textI, texti1, texti2):
    preds = model(textI, texti1, texti2)
    filename = generate_random_filename(length=15, extension='.mp4')
    plot_3d_motion(filename, HML_KINEMATIC_CHAIN, preds, title="", fps=30)
    return filename

model = AutoModel.from_pretrained("pabloruizponce/in2IN", trust_remote_code=True)

demo = gr.Interface(fn=generate, 
                    inputs=[gr.Text(label="Interaction Description"),
                            gr.Text(label="Individual1 Description"),
                            gr.Text(label="Individual2 Description")],
                    outputs=gr.Video(),
                    title="in2IN: Leveraging individual Information to Generate Human INteractions",
                    examples=[
                        ["The two guys meet, grip each other’s hand, and nod in agreement.", "One person extends their hand forward and grips the hand of the other person.", "One person nods their head in agreement."],
                        ["One person spots the other person on the street, lifts the right hand to greet, and the other person glances towards one person.", "One person raises their right hand in greeting.", "One person turns their head to make eye contact."],
                        ["Both they lift their right legs to kick one another.", "One person lifts their right leg in preparation to kick.", "The other person also lifts their right leg in preparation to kick."],
                        ["The two individuals are dancing ballroom together.", "One person is gracefully twirling and gliding across the dance floor.", "The other person leads with confident steps and skillfully guides their partner through intricate dance moves."]
                    ])

demo.launch()