import gradio as gr
from gradio_client import Client, handle_file

def generate_video(input_image, prompt, negative_prompt, diffusion_step, height, width, scfg_scale, use_dctinit, dct_coefficients, noise_level, motion_bucket_id, seed):
    client = Client("maxin-cn/Cinemo")
    
    result = client.predict(
        input_image=handle_file(input_image),
        prompt=prompt,
        negative_prompt=negative_prompt,
        diffusion_step=diffusion_step,
        height=height,
        width=width,
        scfg_scale=scfg_scale,
        use_dctinit=use_dctinit,
        dct_coefficients=dct_coefficients,
        noise_level=noise_level,
        motion_bucket_id=motion_bucket_id,
        seed=seed,
        api_name="/gen_video"
    )
    print("API response" , result)
    
    video_path = result.get('video')  # Extract the video file path from the API response
    
    if video_path is None:
        return "The API did not return a valid video. Please try again."
    
    return video_path  # Return the path to the video file

# Gradio Interface
with gr.Blocks() as demo:
    with gr.Row():
        input_image = gr.Image(label="Input Image", type="filepath")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt")
        negative_prompt = gr.Textbox(label="Negative Prompt")
    with gr.Row():
        diffusion_step = gr.Slider(label="Sampling steps", minimum=1, maximum=100, value=50)
        height = gr.Slider(label="Height", minimum=64, maximum=1024, value=320)
        width = gr.Slider(label="Width", minimum=64, maximum=1024, value=512)
        scfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, value=7.5)
    with gr.Row():
        use_dctinit = gr.Checkbox(label="Enable DCTInit", value=True)
        dct_coefficients = gr.Slider(label="DCT Coefficients", minimum=0.0, maximum=1.0, value=0.23)
        noise_level = gr.Slider(label="Noise Level", minimum=0, maximum=1000, value=985)
        motion_bucket_id = gr.Slider(label="Motion Intensity", minimum=0, maximum=20, value=10)
        seed = gr.Number(label="Seed", value=100)

    video_output = gr.Video(label="Generated Video")
    
    generate_button = gr.Button("Generate Video")
    
    generate_button.click(generate_video, inputs=[input_image, prompt, negative_prompt, diffusion_step, height, width, scfg_scale, use_dctinit, dct_coefficients, noise_level, motion_bucket_id, seed], outputs=video_output)

demo.launch()