'commit'

dsLightRag/JiMeng/JmTxt2Video.py

@@ -0,0 +1,122 @@
+import json
+import time
+import logging
+from JiMeng.Kit.JmCommon import JmCommon
+from JiMeng.Kit.JmErrorCode import JmErrorCode
+
+# 配置日志
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger('JmTxt2Video')
+
+class JmTxt2Video(JmCommon):
+    action = "CVSync2AsyncSubmitTask"
+    req_key = "jimeng_vgfm_t2v_l20"
+
+    @staticmethod
+    def generate_video(prompt):
+        """
+        生成视频
+
+        :param prompt: 提示词
+        :return: 任务结果
+        :raises Exception: 异常信息
+        """
+        req = {
+            "req_key": JmTxt2Video.req_key,
+            "prompt": prompt
+        }
+
+        response_body = JmCommon.do_request(
+            method="POST",
+            query_list={},  # 添加空字典作为query_list参数
+            headers={},
+            body=json.dumps(req).encode('utf-8'),
+            action=JmTxt2Video.action
+        )
+        return json.loads(response_body)
+
+
+def main():
+    try:
+        # 示例提示词
+        # prompt = "蓝色毛绒玩具在超市里拖地，结果拖把洒出好多五颜六色的粉末，接着把粉末洒向镜头前，镜头随之穿过粉末"
+        prompt = "第一个镜头-全景：身穿宇航服的狗狗左右环顾第二个镜头-近景：宇航员狗狗向操作台走去第三个镜头-特写：宇航员狗狗伸手去按操作界"
+        # 保存的文件名称
+        mp4_file_name = "Text2Video.mp4"
+
+        # 添加重试逻辑，处理API并发限制错误
+        jo = None
+        generate_retry_count = 0
+        max_generate_retries = 1000  # 最大重试次数
+        generate_retry_interval = 5000  # 重试间隔（毫秒）
+
+        while generate_retry_count < max_generate_retries:
+            try:
+                jo = JmTxt2Video.generate_video(prompt)
+                logger.info(f"结果: {jo}")
+
+                code = jo.get("code")
+                if code == JmErrorCode.API_CONCURRENT_LIMIT[0]:
+                    logger.warning(f"API并发限制，等待{generate_retry_interval}毫秒后重试...")
+                    time.sleep(generate_retry_interval / 1000)
+                    generate_retry_count += 1
+                    continue
+                elif not JmErrorCode.is_success(code):
+                    logger.error(f"生成视频失败: 错误码={code}, 错误信息={JmErrorCode.get_message_by_code(code)}")
+                    return
+                # 成功获取结果，跳出循环
+                break
+            except Exception as e:
+                logger.error(f"生成视频异常: {str(e)}", exc_info=True)
+                generate_retry_count += 1
+                if generate_retry_count < max_generate_retries:
+                    logger.warning(f"等待{generate_retry_interval}毫秒后重试...")
+                    time.sleep(generate_retry_interval / 1000)
+                else:
+                    raise  # 达到最大重试次数，抛出异常
+
+        if generate_retry_count >= max_generate_retries:
+            logger.error(f"生成视频失败，已达到最大重试次数: {max_generate_retries}")
+            return
+
+        task_id = jo.get("data", {}).get("task_id")
+        logger.info(f"任务ID: {task_id}")
+
+        # 检查任务是不是已经结束
+        retry_count = 0
+        max_retries = 1000  # 最大重试次数
+        retry_interval = 3000  # 重试间隔（毫秒）
+
+        while retry_count < max_retries:
+            result = JmCommon.query_task_result(task_id)
+            logger.info(f"查询结果: {result}")
+
+            code = result.get("code")
+            if not JmErrorCode.is_success(code):
+                logger.error(f"查询失败: 错误码={code}, 错误信息={JmErrorCode.get_message_by_code(code)}")
+                break
+
+            data = result.get("data")
+            if data and data.get("video_url"):
+                video_url = data.get("video_url")
+                logger.info(f"视频地址：{video_url}")
+                # 下载mp4
+                save_video_path = f"{JmCommon.base_path}{mp4_file_name}"
+                logger.info("开始下载视频...")
+                JmCommon.download_file(video_url, save_video_path)
+                logger.info(f"视频已下载到: {save_video_path}")
+                break
+            else:
+                logger.info(f"任务处理中，等待{retry_interval}毫秒后重试...")
+                time.sleep(retry_interval / 1000)
+                retry_count += 1
+
+        if retry_count >= max_retries:
+            logger.error(f"任务查询超时，已达到最大重试次数: {max_retries}")
+
+    except Exception as e:
+        logger.error(f"程序执行异常: {str(e)}", exc_info=True)
+
+
+if __name__ == "__main__":
+    main()