import requests
import time
import gradio as gr
from zhipuai import ZhipuAI

# 请填写您自己的APIKey
api_key = "91dceba38eabfcf8d084fddd9b1cf3da.bCjO9cDi6S5kAy5m"
client = ZhipuAI(api_key=api_key)

# 设置API密钥和请求URL
generate_url = "https://open.bigmodel.cn/api/paas/v4/videos/generations"
result_url_template = "https://open.bigmodel.cn/api/paas/v4/async-result/{id}"

# 请求头
headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {api_key}"
}

def generate_video(prompt):
    if not prompt:
        return "输入不能为空，请输入视频描述。", None

    try:
        # 请求数据
        data = {
            "model": "cogvideox",
            "prompt": prompt
        }

        # 发送POST请求生成视频
        response = requests.post(generate_url, headers=headers, json=data)
        response_data = response.json()
        print("视频生成请求响应:", response_data)

        # 从生成视频的响应中获取任务ID
        task_id = response_data.get("id")
        if not task_id:
            return "无法获取任务ID，请检查请求响应。", None

        # 查询结果的URL
        result_url = result_url_template.format(id=task_id)

        # 轮询查询视频生成结果
        while True:
            result_response = requests.get(result_url, headers=headers)
            result_data = result_response.json()
            print("查询视频生成结果响应:", result_data)
            task_status = result_data.get("task_status")

            if task_status == "SUCCESS":
                print("视频生成成功!")
                video_info = result_data.get("video_result", [])[0]
                video_url = video_info.get("url")
                cover_image_url = video_info.get("cover_image_url")
                return video_url, cover_image_url
            elif task_status == "FAIL":
                print("视频生成失败!")
                return f"视频生成失败! 错误信息: {result_data}", None
            else:
                print("视频生成处理中，请稍候...")
                time.sleep(10)  # 等待10秒后再次查询

    except Exception as e:
        print(f"发生错误: {str(e)}")
        return f"发生错误: {str(e)}", None

# 创建Gradio界面
iface = gr.Interface(
    fn=generate_video,
    inputs=[gr.Textbox(label="输入视频描述 (prompt)", placeholder="输入描述生成视频...", lines=2)],
    outputs=[gr.Video(label="生成的视频 (Video)"), gr.Image(label="视频封面 (Cover Image)")],
    title="金润AI视频生成",
    description="输入描述生成视频，并查看结果。"
)

# 启动Gradio界面
iface.launch()