video-create/tools/image_generator.py

import os
import time
import requests
import json
import threading
import asyncio
import aiohttp
from typing import Optional, Dict, Callable
from dotenv import load_dotenv
from interfaces.image_output import ImageOutput
from utils.tools import encode_image, download_image, download_video
from utils.upload import upload_file_to_tos
from utils.logger_config import setup_logger

load_dotenv()

logger = setup_logger(__name__)

class ArkImageGenerator:
    """Ark 图片生成 API 封装类"""
    
    def __init__(
        self,
        auth_token: str = None,
        model: str = "doubao-seedream-4-0-250828",
        sequential_generation: str = "disabled",
        response_format: str = "url",
        stream: bool = False,
        watermark: bool = True,
        timeout: int = 120
    ):
        """
        初始化图片生成器
        
        参数:
            auth_token: 认证令牌（Bearer Token）
            model: 模型名称（固定配置）
            sequential_generation: 序列生成开关（固定配置）
            response_format: 响应格式（固定配置）
            stream: 流式响应开关（固定配置）
            watermark: 水印开关（固定配置）
            timeout: 请求超时时间（秒）
        """
        self.api_url = "https://ark.cn-beijing.volces.com/api/v3/images/generations"

        if not auth_token:
            auth_token = os.getenv("ARK_API_KEY")

        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {auth_token}"
        }
        # 固定配置参数
        self.config = {
            "model": model,
            "sequential_image_generation": sequential_generation,
            "response_format": response_format,
            "stream": stream,
            "watermark": watermark
        }
        self.timeout = timeout

    async def generate_without_refer(self, prompt: str, size: str = "1440x2560") -> Optional[Dict]:
        if not prompt:
            logger.info("错误：prompt不能为空")
            return None

        payload = {
            **self.config,
            "prompt": prompt,
            "size": size
        }
        try:
            # 使用 aiohttp 进行异步请求
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    url=self.api_url,
                    headers=self.headers,
                    json=payload,
                    timeout=self.timeout
                ) as response:
                    response.raise_for_status()
                    result_data = await response.json()
                    result_image = result_data["data"][0]["url"]
                    return ImageOutput(fmt="url", ext="png", data=result_image)
        
        except Exception as e:
            logger.info(f"错误：生成图片时发生异常：{e}")
            return None

    async def generate(self, prompt: str, image_url: list[str], size: str = "1440x2560") -> Optional[Dict]:
        # 验证必填参数
        if not prompt or not image_url:
            logger.info("错误：prompt和image_url不能为空")
            return None

        # 如果image_url为图片路径，则编码为base64格式
        reference_image = [encode_image(image_url[i]) if "http" not in image_url[i] else image_url[i] for i in range(len(image_url))]
        
        # 构建请求体（合并固定配置和动态参数）
        payload = {
            **self.config,
            "prompt": prompt,
            "image": reference_image,
            "size": size
        }
        try:
            # 使用 aiohttp 进行异步请求
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    url=self.api_url,
                    headers=self.headers,
                    json=payload,
                    timeout=self.timeout
                ) as response:
                    response.raise_for_status()
                    result_data = await response.json()
                    result_image = result_data["data"][0]["url"]
                    return ImageOutput(fmt="url", ext="png", data=result_image)
        
        except Exception as e:
            logger.info(f"错误：生成图片时发生异常：{e}")
            return None

class ArkVideoGenerator:
    """Ark 图生视频 API 封装类，支持通过参考图和文本描述生成视频"""
    
    def __init__(
        self,
        auth_token: str = None,
        model: str = "doubao-seedance-1-0-pro-250528",
        timeout: int = 60,
        poll_interval: int = 5,  # 轮询间隔（秒）
        max_poll_time: int = 500  # 最大轮询总时间（秒）
    ):
        # 固定 API 端点
        self.api_url = "https://ark.cn-beijing.volces.com/api/v3/contents/generations/tasks"

        if not auth_token:
            auth_token = os.getenv("ARK_API_KEY")
        # 固定请求头
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {auth_token}"
        }
        # 固定配置参数（模型、超时）
        self.model = model
        self.timeout = timeout
        self.poll_interval = poll_interval
        self.max_poll_time = max_poll_time

    def create_video_task(
        self,
        prompt: str,
        image_url: str,
        gen_params: str = ""
    ) -> Optional[Dict]:
        # 1. 验证动态参数合法性
        if not prompt.strip():
            logger.info("错误：文本描述（text_prompt）不能为空")
            return None
        if not image_url.strip():
            logger.info("错误：参考图 URL（reference_image_url）不能为空")
            return None

        refernece_image = encode_image(image_url) if "http" not in image_url else image_url
        
        # 2. 构建请求体（按 API 要求格式组装 content 列表）
        payload = {
            "model": self.model,
            "content": [
                {
                    "type": "text",
                    "text": prompt + gen_params
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": refernece_image
                    }
                }
            ]
        }
        
        # 3. 发送 POST 请求并处理响应
        try:
            response = requests.post(
                url=self.api_url,
                headers=self.headers,
                data=json.dumps(payload, ensure_ascii=False),
                timeout=self.timeout
            )
            response.raise_for_status()
            return response.json()
        
        except Exception as e:
            logger.info(f"创建任务失败：{str(e)}")
            return None

    def query_video_task(self, task_id: str) -> Optional[Dict]:
        """
        新增：查询图生视频任务结果
        
        参数:
            task_id: 视频任务 ID（从 create_video_task 响应中获取）
        
        返回:
            任务结果详情（含视频状态、视频 URL 等）或 None
        """
        # 1. 验证任务 ID
        if not task_id.strip():
            logger.info("错误：任务 ID（task_id）不能为空")
            return None
        
        # 2. 构建查询 URL（拼接 task_id）
        query_url = f"{self.api_url}/{task_id}"
        
        # 3. 发送 GET 请求查询结果
        try:
            response = requests.get(
                url=query_url,
                headers=self.headers,
                timeout=self.timeout
            )
            # 触发 HTTP 错误（如 404 任务不存在、401 令牌无效）
            response.raise_for_status()
            return response.json()

        except Exception as e:
            logger.info(f"查询任务失败：{str(e)}")
            return None

    def _background_poll(
        self,
        task_id: str,
        filename: str,
        callback: Callable[[str, str, Optional[Dict], Optional[str]], None]
    ):
        """
        后台轮询任务状态的线程函数
        :param task_id: 任务ID
        :param callback: 回调函数，参数为 (task_id, 成功结果, 错误信息)
        """
        start_time = time.time()
        while True:
            elapsed = time.time() - start_time
            if elapsed > self.max_poll_time:
                callback(task_id, filename, None, f"任务超时（超过 {self.max_poll_time} 秒）")

            # 查询任务状态
            result = self.query_video_task(task_id)
            if not result:
                time.sleep(self.poll_interval)
                continue
            
            # 解析状态
            status = result.get("status", "").lower()
            if status == "succeeded":
                callback(task_id, filename, result, None)
                return 
            elif status == "failed":
                error_msg = result.get("error", {}).get("message", "未知错误")
                callback(task_id, filename, None, error_msg)
                return
            elif status in ["pending", "processing"]:
                logger.info(f"任务 {task_id} 处理中（{int(elapsed)}秒），状态：{status}")
                time.sleep(self.poll_interval)
            else:
                logger.info(f"任务 {task_id} 未知状态：{status}，继续等待...")
                time.sleep(self.poll_interval)

    def create_video_task_async(
        self,
        prompt: str,
        image_url: str,
        gen_params: str,
        filename: str,
        callback: Callable[[str, str, Optional[Dict], Optional[str]], None]
    ) -> Optional[str]:
        # 1. 提交任务
        task_response = self.create_video_task(prompt, image_url, gen_params)
        if not task_response or "id" not in task_response:
            logger.info("任务提交失败，无法启动后台轮询")
            return None

        task_id = task_response["id"]
        logger.info(f"任务提交成功，task_id: {task_id}，启动后台轮询...")

        # 2. 启动后台线程轮询结果
        poll_thread = threading.Thread(
            target=self._background_poll,
            args=(task_id, filename, callback),
            daemon=True  # 守护线程：主程序退出时自动结束
        )
        poll_thread.start()

        return task_id
        
# 1. 定义回调函数：任务完成/失败时会被调用
def handle_video_result(task_id: str, filename, result: Optional[Dict], error: Optional[str]) -> None:
    if error:
        logger.info(f"\n任务 {task_id} 处理失败：{error}")
    else:
        video_url = result.get("content", {}).get("video_url")
        output_path = "./output/" + filename
        download_video(video_url, output_path) 
        logger.info(f"生成视频已下载：{output_path}")
                

# API配置
API_URL = os.getenv("AUDIO_GEN_API")

def audio_generator(text, spk_audio="./data/audio/voice_07.wav", emo_audio="./data/audio/emo_sad.wav"):
    """调用TTS API生成语音"""
    payload = {
        "text": text,
        "spk_audio_prompt": spk_audio,
        "emo_audio_prompt": emo_audio
    }
    
    response = requests.post(API_URL, json=payload)
    
    if response.status_code == 200:
        result = response.json()
        if result["status"] == "success":
            print(f"语音生成成功: {result['audio_file']}")
            return result["audio_file"]
    else:
        print(f"请求失败: {response.text}")
        return None


image_generator = ArkImageGenerator()
video_generator = ArkVideoGenerator()
    
if __name__ == "__main__":
    # 1. 初始化生成器（配置固定参数）
    image_generator = ArkImageGenerator()
    video_generator = ArkVideoGenerator()
    
    # 2. 调用生成方法（仅传入动态参数）
    result = image_generator.generate(
        prompt="狗狗在草地上追逐蒲公英",
        image_url="https://ark-project.tos-cn-beijing.volces.com/doc_image/seedream4_imageToimage.png",
        filename="1.jpg"
    )
    
    logger.info(f"result:{result}")

    # print(video_generator.query_video_task("cgt-20251022103303-9hgrr"))

    # cgt-20251022101137-852fw cgt-20251022102536-kt8pj cgt-20251022103303-9hgrr

    # task_id = video_generator.create_video_task_async(
    #     prompt="狗狗不停地在草地上跳跃",
    #     image_url="https://ark-project.tos-cn-beijing.volces.com/doc_image/seedream4_imageToimage.png",
    #     gen_params="",
    #     filename="1.mp4",
    #     callback=handle_video_result
    # )

    # if task_id:
    #     print("\n主流程：任务已提交，开始执行其他操作...")
    #     for i in range(10):
    #         print(f"主流程：正在执行第 {i+1} 步操作...")
    #         time.sleep(1)  # 模拟主流程耗时操作
    #     print("主流程：所有操作执行完毕，等待后台任务结果（若未完成）...")
        
    #     # 防止主程序提前退出（实际生产环境可能有其他阻塞逻辑）
    #     # 这里仅为演示：等待所有后台线程完成
    #     while threading.active_count() > 1:
    #         time.sleep(1)