dsProject/dsLightRag/Util/VideoRetalk.py

import requests
import time
from typing import Dict, Optional

from Config import Config


class VideoRetalk:
    """
    阿里云DashScope LivePortrait视频生成类
    videoretalk是一个人物视频生成模型，可基于人物视频和人声音频，生成人物讲话口型与输入音频相匹配的新视频。
    视频口型替换-声动人像VideoRetalk
    实现图像和音频合成视频功能
    """
    
    def __init__(self, api_key: str):
        """
        初始化视频生成类
        
        Args:
            api_key: 阿里云DashScope API密钥
        """
        self.api_key = api_key
        self.base_url = "https://dashscope.aliyuncs.com/api/v1"
        self.video_synthesis_url = f"{self.base_url}/services/aigc/image2video/video-synthesis"
    
    def submit_video_task(self, image_url: str, audio_url: str, 
                         template_id: str = "normal", 
                         eye_move_freq: float = 0.5,
                         video_fps: int = 30,
                         mouth_move_strength: float = 1.0,
                         paste_back: bool = True,
                         head_move_strength: float = 0.7) -> Dict:
        """
        提交视频生成任务
        
        Args:
            image_url: 输入图片URL
            audio_url: 输入音频URL
            template_id: 模板ID，默认为"normal"
            eye_move_freq: 眼睛移动频率，默认0.5
            video_fps: 视频帧率，默认30
            mouth_move_strength: 嘴巴移动强度，默认1.0
            paste_back: 是否粘贴背景，默认True
            head_move_strength: 头部移动强度，默认0.7
            
        Returns:
            Dict: 包含task_id和task_status的响应数据
        """
        headers = {
            'X-DashScope-Async': 'enable',
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }
        
        payload = {
            "model": "liveportrait",
            "input": {
                "image_url": image_url,
                "audio_url": audio_url
            },
            "parameters": {
                "template_id": template_id,
                "eye_move_freq": eye_move_freq,
                "video_fps": video_fps,
                "mouth_move_strength": mouth_move_strength,
                "paste_back": paste_back,
                "head_move_strength": head_move_strength
            }
        }
        
        try:
            response = requests.post(self.video_synthesis_url, 
                                   headers=headers, 
                                   json=payload)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            raise Exception(f"提交视频任务失败: {e}")
    
    def get_task_status(self, task_id: str) -> Dict:
        """
        查询任务状态
        
        Args:
            task_id: 任务ID
            
        Returns:
            Dict: 任务状态信息
        """
        headers = {
            'Authorization': f'Bearer {self.api_key}'
        }
        
        task_url = f"{self.base_url}/tasks/{task_id}"
        
        try:
            response = requests.get(task_url, headers=headers)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            raise Exception(f"查询任务状态失败: {e}")
    
    def wait_for_task_completion(self, task_id: str, 
                               poll_interval: int = 5, 
                               timeout: int = 300) -> Dict:
        """
        等待任务完成
        
        Args:
            task_id: 任务ID
            poll_interval: 轮询间隔(秒)，默认5秒
            timeout: 超时时间(秒)，默认300秒
            
        Returns:
            Dict: 任务完成后的结果
        """
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            task_status = self.get_task_status(task_id)
            
            status = task_status.get('output', {}).get('task_status')
            
            if status == 'SUCCEEDED':
                return task_status
            elif status == 'FAILED':
                error_code = task_status.get('output', {}).get('code', '未知错误')
                error_message = task_status.get('output', {}).get('message', '无错误信息')
                raise Exception(f"任务执行失败: {error_code} - {error_message}")
            elif status in ['PENDING', 'RUNNING']:
                print(f"任务状态: {status}, 等待中...")
                time.sleep(poll_interval)
            else:
                raise Exception(f"未知的任务状态: {status}")
        
        raise Exception(f"任务超时，未在{timeout}秒内完成")
    
    def generate_video(self, image_url: str, audio_url: str, 
                      **kwargs) -> Optional[str]:
        """
        生成视频的完整流程
        
        Args:
            image_url: 输入图片URL
            audio_url: 输入音频URL
            **kwargs: 其他参数，同submit_video_task
            
        Returns:
            str: 生成的视频URL，失败返回None
        """
        try:
            # 提交任务
            submit_response = self.submit_video_task(image_url, audio_url, **kwargs)
            task_id = submit_response.get('output', {}).get('task_id')
            
            if not task_id:
                print("提交任务失败，未获取到task_id")
                return None
            
            print(f"任务已提交，task_id: {task_id}")
            
            # 等待任务完成
            result = self.wait_for_task_completion(task_id)
            
            # 获取视频URL
            video_url = result.get('output', {}).get('results', {}).get('video_url')
            
            if video_url:
                print(f"视频生成成功: {video_url}")
                # 获取使用情况信息
                usage = result.get('usage', {})
                video_info = {
                    'video_url': video_url,
                    'video_duration': usage.get('video_duration'),
                    'video_ratio': usage.get('video_ratio')
                }
                if usage:
                    print(f"视频时长: {video_info['video_duration']}秒")
                    print(f"视频比例: {video_info['video_ratio']}")
                return video_info
            else:
                print("未找到生成的视频URL")
                return None
                
        except Exception as e:
            print(f"视频生成失败: {e}")
            return None


# 使用示例
if __name__ == "__main__":
    # 替换为您的实际API密钥
    API_KEY = Config.ALY_LLM_API_KEY
    
    # 创建视频生成实例
    video_retalk = VideoRetalk(API_KEY)
    
    # 示例：生成视频
    try:
        video_url = video_retalk.generate_video(
            image_url="https://dsideal.obs.cn-north-1.myhuaweicloud.com/HuangHai/Backup/LiBai.jpg",
            audio_url="https://dsideal.obs.cn-north-1.myhuaweicloud.com/HuangHai/Backup/JiangJinJiu.mp3",
            template_id="normal",
            eye_move_freq=0.5,
            video_fps=30,
            mouth_move_strength=1.0,
            paste_back=True,
            head_move_strength=0.7
        )
        
        if video_url:
            print(f"最终视频URL: {video_url}")
        
    except Exception as e:
        print(f"执行过程中发生错误: {e}")