Integrate Zhipu AI

Author: zack-zzq

api/zhipu.py

@@ -0,0 +1,65 @@
+import logging
+import time
+
+from zhipuai import ZhipuAI
+
+# 设置日志
+logging.basicConfig(level=logging.WARNING)
+logger = logging.getLogger(__name__)
+
+# 初始化客户端（API Key 通过环境变量传入）
+client = ZhipuAI()
+
+
+def generate_zhipu(prompt, image_url=None, model="cogvideox-2", quality="speed", with_audio=False, size="1920x1080",
+                   fps=30):
+    try:
+        # 根据是否有图片选择文生视频或图生视频
+        params = {
+            "model": model,
+            "prompt": prompt,
+            "quality": quality,
+            "with_audio": with_audio,
+        }
+        if model != "cogvideox-flash":
+            params["size"] = size
+            params["fps"] = fps
+        if image_url:
+            params["image_url"] = image_url
+
+        # 发起生成请求
+        response = client.videos.generations(**params)
+        task_id = response.id
+        logger.debug(f"智谱AI任务创建成功: {task_id}")
+
+        # 轮询任务状态
+        while True:
+            result = client.videos.retrieve_videos_result(id=task_id)
+            logger.debug(f"任务状态: {result.task_status}")
+            if result.task_status in ["SUCCESS", "FAILED"]:
+                break
+            time.sleep(2)
+
+        # 打印完整结果以调试
+        logger.debug(f"任务结果完整内容: {vars(result)}")
+
+        if result.task_status == "SUCCESS":
+            # 处理 video_result
+            video_result = result.video_result
+            if isinstance(video_result, list) and video_result:
+                # 如果是列表，取第一个视频 URL
+                video_url = video_result[0].url
+            elif hasattr(video_result, "url"):
+                # 如果是单一对象，直接取 url 属性
+                video_url = video_result.url
+            else:
+                return None, "任务成功但未找到视频 URL"
+
+            if video_url:
+                return video_url, "任务完成"
+            return None, "任务成功但未找到视频 URL"
+        else:
+            return None, f"任务失败: {result.task_status}"
+    except Exception as e:
+        logger.error(f"智谱AI生成失败: {str(e)}")
+        return None, f"智谱AI生成失败: {str(e)}"

main.py

@@ -7,64 +7,58 @@
 from api.ark import generate_volcengine
 from api.bailian import generate_aliyun, MODEL_MAPPING
 from api.tebi import upload_file_to_tebi
+from api.zhipu import generate_zhipu
 
-# 设置全局日志级别为 WARNING，减少详细输出
+# 设置全局日志级别为 WARNING
 logging.basicConfig(level=logging.WARNING)
 logging.getLogger("httpx").setLevel(logging.WARNING)
 
 
 # 检查图片是否符合火山引擎要求
 def validate_image(file_path):
     try:
-        # 文件大小检查（小于 10MB）
         file_size = os.path.getsize(file_path)
-        if file_size > 10 * 1024 * 1024:  # 10MB = 10,485,760 字节
+        if file_size > 10 * 1024 * 1024:
             return False, "图片文件大小超过 10MB"
 
-        # 打开图片并检查格式、尺寸
         with Image.open(file_path) as img:
-            # 格式检查
             valid_formats = {"JPEG", "PNG", "WEBP", "BMP", "TIFF"}
             if img.format not in valid_formats:
                 return False, f"图片格式不支持，仅支持 {', '.join(valid_formats)}"
-
-            # 宽高检查
             width, height = img.size
             aspect_ratio = width / height
             min_side = min(width, height)
             max_side = max(width, height)
-
             if not (0.4 <= aspect_ratio <= 2.5):
                 return False, "图片宽高比需在 2:5 到 5:2 之间 (0.4 - 2.5)"
             if min_side < 300:
                 return False, "图片短边像素需大于等于 300px"
             if max_side > 6000:
                 return False, "图片长边像素需小于等于 6000px"
-
         return True, "图片验证通过"
     except Exception as e:
         return False, f"图片验证失败: {str(e)}"
 
 
 # 判断是否需要图片上传
-def is_image_required(platform, aliyun_model, ark_duration):
+def is_image_required(platform, aliyun_model, ark_duration, zhipu_model):
     if platform == "火山引擎":
-        return ark_duration == 5  # 仅 5 秒支持文图生视频
+        return ark_duration == 5
     elif platform == "阿里云百炼":
         return aliyun_model in ["通义万相-图生视频2.1-Turbo", "通义万相-图生视频2.1-Plus"]
+    elif platform == "智谱AI":
+        return zhipu_model == "CogVideoX-2 (图生视频)"
     return False
 
 
 # 生成视频逻辑
-def generate_video(platform, prompt, image_file, aliyun_model, ark_ratio, ark_duration, bailian_size):
+def generate_video(platform, prompt, image_file, aliyun_model, ark_ratio, ark_duration, bailian_size, zhipu_model,
+                   zhipu_quality, zhipu_audio, zhipu_size, zhipu_fps):
     image_url = None
-    if is_image_required(platform, aliyun_model, ark_duration) and image_file:
-        # 验证图片
+    if is_image_required(platform, aliyun_model, ark_duration, zhipu_model) and image_file:
         is_valid, message = validate_image(image_file)
         if not is_valid:
             return None, message
-
-        # 上传到 Tebi
         image_url = upload_file_to_tebi(image_file)
         if not image_url:
             return None, "图片上传到 Tebi 失败"
@@ -73,15 +67,18 @@ def generate_video(platform, prompt, image_file, aliyun_model, ark_ratio, ark_du
         video_url, status = generate_volcengine(prompt, image_url, ark_ratio, ark_duration)
     elif platform == "阿里云百炼":
         video_url, status = generate_aliyun(prompt, image_url, aliyun_model, bailian_size)
+    elif platform == "智谱AI":
+        model = "cogvideox-2" if "CogVideoX-2" in zhipu_model else "cogvideox-flash"
+        video_url, status = generate_zhipu(prompt, image_url, model, zhipu_quality, zhipu_audio, zhipu_size, zhipu_fps)
     else:
         return None, "请选择有效平台"
 
     return video_url, status
 
 
 # 更新图片上传区域可见性
-def update_image_visibility(platform, aliyun_model, ark_duration):
-    return gr.update(visible=is_image_required(platform, aliyun_model, ark_duration))
+def update_image_visibility(platform, aliyun_model, ark_duration, zhipu_model):
+    return gr.update(visible=is_image_required(platform, aliyun_model, ark_duration, zhipu_model))
 
 
 # Gradio 界面
@@ -90,10 +87,9 @@ def update_image_visibility(platform, aliyun_model, ark_duration):
     gr.Markdown("输入提示词并选择平台生成视频，支持上传图片用于文图生视频。")
 
     with gr.Row():
-        # 左侧输入区
         with gr.Column(scale=1):
             platform = gr.Dropdown(
-                choices=["火山引擎", "阿里云百炼"],
+                choices=["火山引擎", "阿里云百炼", "智谱AI"],
                 label="选择平台",
                 value="火山引擎"
             )
@@ -103,6 +99,12 @@ def update_image_visibility(platform, aliyun_model, ark_duration):
                 value="通义万相-文生视频2.1-Turbo",
                 visible=False
             )
+            zhipu_model = gr.Dropdown(
+                choices=["CogVideoX-2 (文生视频)", "CogVideoX-2 (图生视频)", "CogVideoX-Flash"],
+                label="智谱AI模型（仅智谱AI生效）",
+                value="CogVideoX-2 (文生视频)",
+                visible=False
+            )
             prompt = gr.Textbox(label="提示词", placeholder="请输入生成视频的描述，例如：一只猫在草地上奔跑")
             image_file = gr.File(label="上传图片（用于文图生视频）", type="filepath")
 
@@ -127,11 +129,30 @@ def update_image_visibility(platform, aliyun_model, ark_duration):
                     value="1280*720"
                 )
 
+            # 智谱AI参数
+            with gr.Group(visible=False) as zhipu_params:
+                zhipu_quality = gr.Dropdown(
+                    choices=["speed", "quality"],
+                    label="输出模式",
+                    value="speed"
+                )
+                zhipu_audio = gr.Checkbox(label="生成AI音效", value=False)
+                zhipu_size = gr.Dropdown(
+                    choices=["720x480", "1024x1024", "1280x960", "960x1280", "1920x1080", "1080x1920", "2048x1080",
+                             "3840x2160"],
+                    label="分辨率",
+                    value="1920x1080"
+                )
+                zhipu_fps = gr.Dropdown(
+                    choices=[30, 60],
+                    label="帧率 (FPS)",
+                    value=30
+                )
+
             with gr.Row():
                 submit_btn = gr.Button("生成视频")
                 clear_btn = gr.Button("清除")
 
-        # 右侧输出区
         with gr.Column(scale=1):
             video_output = gr.Video(label="生成结果")
             status_output = gr.Textbox(label="状态")
@@ -141,36 +162,57 @@ def update_image_visibility(platform, aliyun_model, ark_duration):
     def update_visibility(platform):
         return (
             gr.update(visible=platform == "阿里云百炼"),  # aliyun_model
+            gr.update(visible=platform == "智谱AI"),  # zhipu_model
             gr.update(visible=platform == "火山引擎"),  # ark_params
-            gr.update(visible=platform == "阿里云百炼")  # bailian_params
+            gr.update(visible=platform == "阿里云百炼"),  # bailian_params
+            gr.update(visible=platform == "智谱AI"),  # zhipu_params
+            gr.update(visible=platform == "智谱AI" and "cogvideox-flash" not in zhipu_model.value),  # zhipu_quality
+            gr.update(visible=platform == "智谱AI" and "cogvideox-flash" not in zhipu_model.value),  # zhipu_size
+            gr.update(visible=platform == "智谱AI" and "cogvideox-flash" not in zhipu_model.value)  # zhipu_fps
         )
 
 
     platform.change(
         fn=update_visibility,
         inputs=platform,
-        outputs=[aliyun_model, ark_params, bailian_params]
+        outputs=[aliyun_model, zhipu_model, ark_params, bailian_params, zhipu_params, zhipu_quality, zhipu_size,
+                 zhipu_fps]
     )
     platform.change(
         fn=update_image_visibility,
-        inputs=[platform, aliyun_model, ark_duration],
+        inputs=[platform, aliyun_model, ark_duration, zhipu_model],
         outputs=image_file
     )
     aliyun_model.change(
         fn=update_image_visibility,
-        inputs=[platform, aliyun_model, ark_duration],
+        inputs=[platform, aliyun_model, ark_duration, zhipu_model],
         outputs=image_file
     )
     ark_duration.change(
         fn=update_image_visibility,
-        inputs=[platform, aliyun_model, ark_duration],
+        inputs=[platform, aliyun_model, ark_duration, zhipu_model],
         outputs=image_file
     )
+    zhipu_model.change(
+        fn=update_image_visibility,
+        inputs=[platform, aliyun_model, ark_duration, zhipu_model],
+        outputs=image_file
+    )
+    zhipu_model.change(
+        fn=lambda zhipu_model: (
+            gr.update(visible="cogvideox-flash" not in zhipu_model),
+            gr.update(visible="cogvideox-flash" not in zhipu_model),
+            gr.update(visible="cogvideox-flash" not in zhipu_model)
+        ),
+        inputs=zhipu_model,
+        outputs=[zhipu_quality, zhipu_size, zhipu_fps]
+    )
 
     # 提交生成
     submit_btn.click(
         fn=generate_video,
-        inputs=[platform, prompt, image_file, aliyun_model, ark_ratio, ark_duration, bailian_size],
+        inputs=[platform, prompt, image_file, aliyun_model, ark_ratio, ark_duration, bailian_size, zhipu_model,
+                zhipu_quality, zhipu_audio, zhipu_size, zhipu_fps],
         outputs=[video_output, status_output]
     )
 
@@ -181,4 +223,4 @@ def update_visibility(platform):
         outputs=[video_output, status_output, image_file, prompt]
     )
 
-demo.launch(server_name="0.0.0.0", quiet=True)
+demo.launch(server_name="0.0.0.0")