diff --git a/src/agentscope_bricks/components/__init__.py b/src/agentscope_bricks/components/__init__.py index 333ba7e..eb6e51c 100644 --- a/src/agentscope_bricks/components/__init__.py +++ b/src/agentscope_bricks/components/__init__.py @@ -1,148 +1,185 @@ -# -*- coding: utf-8 -*- -from typing import Dict, Type, List - -from pydantic import BaseModel, Field - -from agentscope_bricks.base import Component -from agentscope_bricks.components.generations.async_image_to_video_wan25 import ( # noqa - ImageToVideoWan25Fetch, - ImageToVideoWan25Submit, -) -from agentscope_bricks.components.generations.async_text_to_video_wan25 import ( # noqa - TextToVideoWan25Submit, - TextToVideoWan25Fetch, -) -from agentscope_bricks.components.generations.image_edit_wan25 import ( - ImageEditWan25, -) -from agentscope_bricks.components.generations.multichannel_speech_to_text import ( # noqa - MultichannelSpeechToText, -) -from agentscope_bricks.components.generations.qwen_image_edit import ( - QwenImageEdit, -) -from agentscope_bricks.components.generations.qwen_image_generation import ( - QwenImageGen, -) -from agentscope_bricks.components.generations.qwen_text_to_speech import ( - QwenTextToSpeech, -) -from agentscope_bricks.components.generations.text_to_video import TextToVideo -from agentscope_bricks.components.generations.image_to_video import ( - ImageToVideo, -) -from agentscope_bricks.components.generations.speech_to_video import ( - SpeechToVideo, -) -from agentscope_bricks.components.searches.modelstudio_search_lite import ( - ModelstudioSearchLite, -) -from agentscope_bricks.components.generations.image_generation import ( - ImageGeneration, -) -from agentscope_bricks.components.generations.image_generation_wan25 import ( - ImageGenerationWan25, -) -from agentscope_bricks.components.generations.image_edit import ImageEdit -from agentscope_bricks.components.generations.image_style_repaint import ( - ImageStyleRepaint, -) -from agentscope_bricks.components.generations.speech_to_text import ( - SpeechToText, -) - -from agentscope_bricks.components.generations.async_text_to_video import ( - TextToVideoSubmit, - TextToVideoFetch, -) -from agentscope_bricks.components.generations.async_image_to_video import ( - ImageToVideoSubmit, - ImageToVideoFetch, -) -from agentscope_bricks.components.generations.async_speech_to_video import ( - SpeechToVideoSubmit, - SpeechToVideoFetch, -) -from agentscope_bricks.components.generations.async_text_to_video_wan26 import ( # noqa - TextToVideoWan26Submit, -) -from agentscope_bricks.components.generations.async_image_to_video_wan26 import ( # noqa - ImageToVideoWan26Submit, -) -from agentscope_bricks.components.generations.image_generation_wan26 import ( # noqa - ImageGenerationWan26, -) -from agentscope_bricks.components.generations.fetch_wan import WanVideoFetch -from agentscope_bricks.components.generations.qwen_image_edit_new import ( - QwenImageEditNew, -) # noqa - - -class McpServerMeta(BaseModel): - instructions: str = Field( - ..., - description="服务描述", - ) - components: List[Type[Component]] = Field( - ..., - description="组件列表", - ) - - -mcp_server_metas: Dict[str, McpServerMeta] = { - "modelstudio_wan_image": McpServerMeta( - instructions="基于通义万相大模型的智能图像生成服务，提供高质量的图像处理和编辑功能", - components=[ImageGeneration, ImageEdit, ImageStyleRepaint], - ), - "modelstudio_wan_video": McpServerMeta( - instructions="基于通义万相大模型提供AI视频生成服务，支持文本到视频、图像到视频和语音到视频的多模态生成功能", - components=[ - TextToVideoSubmit, - TextToVideoFetch, - ImageToVideoSubmit, - ImageToVideoFetch, - SpeechToVideoSubmit, - SpeechToVideoFetch, - ], - ), - "modelstudio_wan25_media": McpServerMeta( - instructions="基于通义万相大模型2.5版本提供的图像和视频生成服务", - components=[ - ImageGenerationWan25, - ImageEditWan25, - TextToVideoWan25Submit, - TextToVideoWan25Fetch, - ImageToVideoWan25Submit, - ImageToVideoWan25Fetch, - ], - ), - "modelstudio_qwen_image": McpServerMeta( - instructions="基于通义千问大模型的智能图像生成服务，提供高质量的图像处理和编辑功能", - components=[ - QwenImageGen, - QwenImageEdit, - QwenImageEditNew, - ], # ← 新增 QwenImageEditNew - ), - "modelstudio_web_search": McpServerMeta( - instructions="提供实时互联网搜索服务，提供准确及时的信息检索功能", - components=[ModelstudioSearchLite], - ), - "modelstudio_speech_to_text": McpServerMeta( - instructions="录音文件的语音识别服务，支持多种音频格式的语音转文字功能", - components=[SpeechToText, MultichannelSpeechToText], - ), - "modelstudio_qwen_text_to_speech": McpServerMeta( - instructions="基于通义千问大模型的语音合成服务，支持多种语言语音合成功能", - components=[QwenTextToSpeech], - ), - "modelstudio_wan26_media": McpServerMeta( - instructions="基于通义万相大模型2.6版本提供的图像和视频生成服务", - components=[ - ImageGenerationWan26, - TextToVideoWan26Submit, - ImageToVideoWan26Submit, - WanVideoFetch, - ], - ), -} +# -*- coding: utf-8 -*- +from typing import Dict, Type, List + +from pydantic import BaseModel, Field + +from agentscope_bricks.base import Component +from agentscope_bricks.components.generations.async_image_to_video_wan25 import ( # noqa + ImageToVideoWan25Fetch, + ImageToVideoWan25Submit, +) +from agentscope_bricks.components.generations.async_text_to_video_wan25 import ( # noqa + TextToVideoWan25Submit, + TextToVideoWan25Fetch, +) +from agentscope_bricks.components.generations.image_edit_wan25 import ( + ImageEditWan25, +) +from agentscope_bricks.components.generations.multichannel_speech_to_text import ( # noqa + MultichannelSpeechToText, +) +from agentscope_bricks.components.generations.qwen_image_edit import ( + QwenImageEdit, +) +from agentscope_bricks.components.generations.qwen_image_generation import ( + QwenImageGen, +) +from agentscope_bricks.components.generations.qwen_text_to_speech import ( + QwenTextToSpeech, +) +from agentscope_bricks.components.generations.text_to_video import TextToVideo +from agentscope_bricks.components.generations.image_to_video import ( + ImageToVideo, +) +from agentscope_bricks.components.generations.speech_to_video import ( + SpeechToVideo, +) +from agentscope_bricks.components.searches.modelstudio_search_lite import ( + ModelstudioSearchLite, +) +from agentscope_bricks.components.generations.image_generation import ( + ImageGeneration, +) +from agentscope_bricks.components.generations.image_generation_wan25 import ( + ImageGenerationWan25, +) +from agentscope_bricks.components.generations.image_edit import ImageEdit +from agentscope_bricks.components.generations.image_style_repaint import ( + ImageStyleRepaint, +) +from agentscope_bricks.components.generations.speech_to_text import ( + SpeechToText, +) + +from agentscope_bricks.components.generations.async_text_to_video import ( + TextToVideoSubmit, + TextToVideoFetch, +) +from agentscope_bricks.components.generations.async_image_to_video import ( + ImageToVideoSubmit, + ImageToVideoFetch, +) +from agentscope_bricks.components.generations.async_speech_to_video import ( + SpeechToVideoSubmit, + SpeechToVideoFetch, +) +from agentscope_bricks.components.generations.async_text_to_video_wan26 import ( # noqa + TextToVideoWan26Submit, +) +from agentscope_bricks.components.generations.async_image_to_video_wan26 import ( # noqa + ImageToVideoWan26Submit, +) +from agentscope_bricks.components.generations.image_generation_wan26 import ( # noqa + ImageGenerationWan26, +) +from agentscope_bricks.components.generations.fetch_wan import WanVideoFetch +from agentscope_bricks.components.generations.qwen_image_edit_new import ( + QwenImageEditNew, +) +from agentscope_bricks.components.generations.image_edit_wan26 import ( + ImageEditWan26, +) +from agentscope_bricks.components.generations.image_generation_zimage import ( + ZImageGeneration, +) +from agentscope_bricks.components.generations.async_image_out_painting import ( + ImageOutPaintingSubmit, + ImageOutPaintingFetch, +) +from agentscope_bricks.components.generations.async_image_to_video_fl_wan22 import ( # noqa + ImageToVideoByFirstAndLastFrameWan22Submit, +) +from agentscope_bricks.components.generations.image_out_painting import ( + ImageOutPaintingAuto, +) +from agentscope_bricks.components.generations.image_text_interleave_generation_wan26 import ( # noqa + WanImageInterleaveGeneration, +) + + +class McpServerMeta(BaseModel): + instructions: str = Field( + ..., + description="服务描述", + ) + components: List[Type[Component]] = Field( + ..., + description="组件列表", + ) + + +mcp_server_metas: Dict[str, McpServerMeta] = { + "modelstudio_wan_image": McpServerMeta( + instructions="基于通义万相大模型的智能图像生成服务，提供高质量的图像处理和编辑功能", + components=[ + ImageGeneration, + ImageEdit, + ImageStyleRepaint, + ImageOutPaintingSubmit, + ImageOutPaintingFetch, + ImageOutPaintingAuto, + ], + ), + "modelstudio_wan_video": McpServerMeta( + instructions="基于通义万相大模型提供AI视频生成服务，支持文本到视频、图像到视频和语音到视频的多模态生成功能", + components=[ + TextToVideoSubmit, + TextToVideoFetch, + ImageToVideoSubmit, + ImageToVideoFetch, + SpeechToVideoSubmit, + SpeechToVideoFetch, + ImageToVideoByFirstAndLastFrameWan22Submit, + WanVideoFetch, + ], + ), + "modelstudio_wan25_media": McpServerMeta( + instructions="基于通义万相大模型2.5版本提供的图像和视频生成服务", + components=[ + ImageGenerationWan25, + ImageEditWan25, + TextToVideoWan25Submit, + TextToVideoWan25Fetch, + ImageToVideoWan25Submit, + ImageToVideoWan25Fetch, + ], + ), + "modelstudio_qwen_image": McpServerMeta( + instructions="基于通义千问大模型的智能图像生成服务，提供高质量的图像处理和编辑功能", + components=[ + QwenImageGen, + QwenImageEdit, + QwenImageEditNew, + ], # ← 新增 QwenImageEditNew + ), + "modelstudio_web_search": McpServerMeta( + instructions="提供实时互联网搜索服务，提供准确及时的信息检索功能", + components=[ModelstudioSearchLite], + ), + "modelstudio_speech_to_text": McpServerMeta( + instructions="录音文件的语音识别服务，支持多种音频格式的语音转文字功能", + components=[SpeechToText, MultichannelSpeechToText], + ), + "modelstudio_qwen_text_to_speech": McpServerMeta( + instructions="基于通义千问大模型的语音合成服务，支持多种语言语音合成功能", + components=[QwenTextToSpeech], + ), + "modelstudio_wan26_media": McpServerMeta( + instructions="基于通义万相大模型2.6版本提供的图像和视频生成服务", + components=[ + ImageGenerationWan26, + TextToVideoWan26Submit, + ImageToVideoWan26Submit, + WanVideoFetch, + ImageEditWan26, + WanImageInterleaveGeneration, + ], + ), + "modelstudio_z_image": McpServerMeta( + instructions="基于通义Z-Image大模型的智能图像生成服务，是一款轻量级文生图模型，" + "可快速生成图像，支持中英文字渲染，并灵活适配多种分辨率与宽高比例。", + components=[ + ZImageGeneration, + ], + ), +} diff --git a/src/agentscope_bricks/components/generations/async_image_out_painting.py b/src/agentscope_bricks/components/generations/async_image_out_painting.py new file mode 100644 index 0000000..98f35e5 --- /dev/null +++ b/src/agentscope_bricks/components/generations/async_image_out_painting.py @@ -0,0 +1,332 @@ +# -*- coding: utf-8 -*- +import os +import uuid +from http import HTTPStatus +from typing import Any, Optional, Dict + +import aiohttp +from mcp.server.fastmcp import Context +from pydantic import BaseModel, Field + +from agentscope_bricks.base.component import Component +from agentscope_bricks.utils.tracing_utils.wrapper import trace +from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key +from agentscope_bricks.utils.tracing_utils import TracingUtil + + +DASHSCOPE_API_BASE = "https://dashscope.aliyuncs.com/api/v1" + + +class ImageOutPaintingSubmitInput(BaseModel): + """ + Input model for submitting an image out-painting (expansion) task. + """ + + image_url: str = Field( + ..., + description="输入图像的公网可访问 URL。", + ) + angle: Optional[float] = Field( + default=None, + description="逆时针旋转角度，取值范围 [0, 359]。默认为 0（不旋转）。", + ) + output_ratio: Optional[str] = Field( + default=None, + description='目标宽高比，可选值：["", "1:1", "3:4", "4:3", "9:16", "16:9"]。' + '默认值为""，表示不设置输出图像的宽高比。', + ) + x_scale: Optional[float] = Field( + default=None, + description="水平方向扩展比例（居中扩展），默认 1.0。可以与 y_scale 搭配使用。取值范围 [1.0, 3.0]。" + "例如：输入图像分辨率为1000×1000（宽×高），x_scale=2.0，扩展后的图像分辨率为2000×1000（宽×高）。" + "保持高度不变，左右各添加500个像素。", + ) + y_scale: Optional[float] = Field( + default=None, + description="垂直方向扩展比例（居中扩展），默认 1.0。可以选择与 x_scale 搭配使用。取值范围 [1.0, 3.0]。" + "例如：输入图像分辨率为1000×1000（宽×高），y_scale=2.0，扩展后的图像分辨率为1000×2000（宽×高）。" + "保持宽度不变，上下各添加500个像素。", + ) + top_offset: Optional[float] = Field( + default=None, + description="在图像上方添加的像素数。默认值为0," + "需满足 top_offset + bottom_offset < 3 × 原图高度。" + "输入图像分辨率为1000×1000（宽×高），top_offset=500，扩展后的图像分辨率为1000×1500（宽×高）。" + "保持宽度不变，只在图像上方添加500个像素。", + ) + bottom_offset: Optional[float] = Field( + default=None, + description="在图像下方添加的像素数。默认值为0，" + "需满足 top_offset + bottom_offset < 3 × 原图高度。" + "例如：输入图像分辨率为1000×1000（宽×高），bottom_offset=500，扩展后的图像分辨率为1000×1500（宽×高）。" + "保持宽度不变，只在图像下方添加500个像素。", + ) + left_offset: Optional[float] = Field( + default=None, + description="在图像左侧添加的像素数。默认值为0，" + "需满足 left_offset + right_offset < 3 × 原图宽度。" + "例如：输入图像分辨率为1000×1000（宽×高），left_offset=500，扩展后的图像分辨率为1500×1000（宽×高）。" + "保持高度不变，只在图像左侧添加500个像素。", + ) + right_offset: Optional[float] = Field( + default=None, + description="在图像右侧添加的像素数。默认值为0，" + "需满足 left_offset + right_offset < 3 × 原图宽度。" + "例如：输入图像分辨率为1000×1000（宽×高），right_offset=500，扩展后的图像分辨率为1500×1000（宽×高）。" + "保持高度不变，只在图像右侧添加500个像素。", + ) + best_quality: Optional[bool] = Field( + default=None, + description="是否启用最佳质量模式。默认 false（速度优先），设为 true 可提升细节但耗时增加。", + ) + limit_image_size: Optional[bool] = Field( + default=None, + description="是否限制输出图像大小（≤5MB）。默认 true，建议保持开启。" + "模型生成的图像需要经过一层安全过滤后才能输出，当前不支持大于10M的图像处理。", + ) + add_watermark: Optional[bool] = Field( + default=None, + description="是否添加水印,True：默认值,添加水印,False：不添加水印。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context containing " + "headers for mcp only, don't generate it", + ) + + +class ImageOutPaintingSubmitOutput(BaseModel): + task_id: str = Field( + title="Task ID", + description="异步任务的唯一标识符，有效期 24 小时。", + ) + task_status: str = Field( + title="Task Status", + description="任务状态：PENDING（排队中）、RUNNING（处理中）、" + "SUCCEEDED（成功）、FAILED（失败）等。", + ) + request_id: Optional[str] = Field( + default=None, + title="Request ID", + description="请求唯一 ID，用于日志追踪。", + ) + + +class ImageOutPaintingSubmit( + Component[ImageOutPaintingSubmitInput, ImageOutPaintingSubmitOutput], +): + name: str = "modelstudio_image_out_painting_submit" + description: str = ( + "图像画面扩展（扩图）异步任务提交工具，基于image-out-painting 模型。\n" + "支持三种扩图方式（按优先级）：\n" + "1. 按宽高比（output_ratio）\n" + "2. 按比例缩放（x_scale / y_scale）\n" + "3. 指定方向像素填充（top/bottom/left/right_offset）\n" + "可选旋转（angle），先旋转后扩图。" + ) + + @trace(trace_type="AIGC", trace_name="image_out_painting_submit") + async def arun( + self, + args: ImageOutPaintingSubmitInput, + **kwargs: Any, + ) -> ImageOutPaintingSubmitOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") + + # 构建 parameters 字典（只包含非 None 值） + parameters: Dict[str, Any] = {} + if args.angle is not None: + parameters["angle"] = args.angle + if args.output_ratio is not None: + parameters["output_ratio"] = args.output_ratio + if args.x_scale is not None: + parameters["x_scale"] = args.x_scale + if args.y_scale is not None: + parameters["y_scale"] = args.y_scale + if args.top_offset is not None: + parameters["top_offset"] = args.top_offset + if args.bottom_offset is not None: + parameters["bottom_offset"] = args.bottom_offset + if args.left_offset is not None: + parameters["left_offset"] = args.left_offset + if args.right_offset is not None: + parameters["right_offset"] = args.right_offset + if args.best_quality is not None: + parameters["best_quality"] = args.best_quality + if args.limit_image_size is not None: + parameters["limit_image_size"] = args.limit_image_size + if args.add_watermark is not None: + parameters["add_watermark"] = args.add_watermark + + headers = { + "Authorization": f"Bearer {api_key}", + "X-DashScope-Async": "enable", + "Content-Type": "application/json", + } + + payload = { + "model": "image-out-painting", + "input": {"image_url": args.image_url}, + "parameters": parameters, + } + + async with aiohttp.ClientSession() as session: + async with session.post( + f"{DASHSCOPE_API_BASE}/services/aigc/image2image/out-painting", + headers=headers, + json=payload, + ) as resp: + status_code = resp.status + response_json = await resp.json() + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "submit_response", + "payload": { + "request_id": request_id, + "response": response_json, + "status_code": status_code, + }, + }, + ) + + if status_code != HTTPStatus.OK or "output" not in response_json: + error_msg = response_json.get("message", "Unknown error") + raise RuntimeError( + f"Failed to submit out-painting task: {error_msg} (code: {status_code})", # noqa + ) + + output = response_json["output"] + task_id = output["task_id"] + task_status = output["task_status"] + actual_request_id = ( + response_json.get("request_id") or request_id or str(uuid.uuid4()) + ) + + return ImageOutPaintingSubmitOutput( + task_id=task_id, + task_status=task_status, + request_id=actual_request_id, + ) + + +# ==================== Fetch Result ==================== + + +class ImageOutPaintingFetchInput(BaseModel): + task_id: str = Field( + ..., + description="要查询的扩图任务 ID。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context containing " + "headers for mcp only, don't generate it", + ) + + +class ImageOutPaintingFetchOutput(BaseModel): + output_image_url: str = Field( + ..., + description="扩图后生成的图像公网 URL（PNG/JPG 等格式）。", + ) + task_id: str = Field( + ..., + description="任务 ID，与输入一致。", + ) + task_status: str = Field( + ..., + description="任务最终状态，成功时为 SUCCEEDED。", + ) + request_id: Optional[str] = Field( + default=None, + description="请求 ID，用于追踪。", + ) + + +class ImageOutPaintingFetch( + Component[ImageOutPaintingFetchInput, ImageOutPaintingFetchOutput], +): + name: str = "modelstudio_image_out_painting_fetch" + description: str = ( + "查询图像画面扩展（扩图）任务的结果。\n" + "输入 Task ID，返回扩图后的图像 URL 和任务状态。\n" + "请在提交任务后轮询此接口，直到状态变为 SUCCEEDED。" + ) + + @trace(trace_type="AIGC", trace_name="image_out_painting_fetch") + async def arun( + self, + args: ImageOutPaintingFetchInput, + **kwargs: Any, + ) -> ImageOutPaintingFetchOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError as e: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") from e + + headers = { + "Authorization": f"Bearer {api_key}", + } + + async with aiohttp.ClientSession() as session: + async with session.get( + f"{DASHSCOPE_API_BASE}/tasks/{args.task_id}", + headers=headers, + ) as resp: + status_code = resp.status + response_json = await resp.json() + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "fetch_response", + "payload": { + "request_id": request_id, + "response": response_json, + "status_code": status_code, + }, + }, + ) + + if status_code != HTTPStatus.OK or "output" not in response_json: + error_msg = response_json.get("message", "Unknown error") + raise RuntimeError( + f"Failed to fetch out-painting result: {error_msg} (code: {status_code})", # noqa + ) + + output = response_json["output"] + task_status = output["task_status"] + + if task_status in ["FAILED", "CANCELED"]: + error_msg = output.get("message", "Task failed") + raise RuntimeError(f"Out-painting task failed: {error_msg}") + + if task_status != "SUCCEEDED": + raise RuntimeError( + f"Task not completed yet. Current status: {task_status}", + ) + + output_image_url = output["output_image_url"] + actual_request_id = ( + response_json.get("request_id") or request_id or str(uuid.uuid4()) + ) + + return ImageOutPaintingFetchOutput( + output_image_url=output_image_url, + task_id=output["task_id"], + task_status=task_status, + request_id=actual_request_id, + ) diff --git a/src/agentscope_bricks/components/generations/async_image_to_video_fl_wan22.py b/src/agentscope_bricks/components/generations/async_image_to_video_fl_wan22.py new file mode 100644 index 0000000..7bd2477 --- /dev/null +++ b/src/agentscope_bricks/components/generations/async_image_to_video_fl_wan22.py @@ -0,0 +1,291 @@ +# -*- coding: utf-8 -*- +import os +import uuid +from http import HTTPStatus +from typing import Any, Optional + +from dashscope.aigc.video_synthesis import AioVideoSynthesis +from mcp.server.fastmcp import Context +from pydantic import BaseModel, Field + +from agentscope_bricks.base.component import Component +from agentscope_bricks.utils.tracing_utils.wrapper import trace +from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key +from agentscope_bricks.utils.tracing_utils import TracingUtil + + +class ImageToVideoByFirstAndLastFrameWan22SubmitInput(BaseModel): + """ + Input model for submitting a + keyframe-to-video task using wan2.2-kf2v-flash. + """ + + first_frame_url: str = Field( + ..., + description="首帧图像，支持公网URL、Base64编码。", + ) + last_frame_url: str = Field( + ..., + description="尾帧图像，支持公网URL、Base64编码。", + ) + prompt: Optional[str] = Field( + default=None, + description="正向提示词，描述希望视频中发生的动作或变化，例如“镜头缓慢推进，风吹动树叶”。", + ) + negative_prompt: Optional[str] = Field( + default=None, + description="反向提示词，用于排除不希望出现的内容，例如“模糊、闪烁、变形、水印”。", + ) + resolution: Optional[str] = Field( + default=None, + description="视频分辨率，可选值：'480P'、'720P'、'1080P'。默认为 '720P'。", + ) + template: Optional[str] = Field( + default=None, + description="不同模型支持不同的特效模板。调用前请查阅视频特效列表，以免调用失败。", + ) + prompt_extend: Optional[bool] = Field( + default=None, + description="Prompt 智能改写。开启后可提升生成效果。默认值为 true。", + ) + watermark: Optional[bool] = Field( + default=None, + description="是否添加水印。false（默认）：不添加；true：添加。", + ) + seed: Optional[int] = Field( + default=None, + description="随机种子，取值范围 [0, 2147483647]。用于提升结果可复现性，但不保证完全一致。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context containing " + "headers for mcp only, don't generate it", + ) + + +class ImageToVideoByFirstAndLastFrameWan22SubmitOutput(BaseModel): + """ + Output of the keyframe-to-video task submission. + """ + + task_id: str = Field( + title="Task ID", + description="异步任务的唯一标识符。", + ) + task_status: str = Field( + title="Task Status", + description="视频生成的任务状态，PENDING：任务排队中，RUNNING：任务处理中，SUCCEEDED：任务执行成功，" + "FAILED：任务执行失败，CANCELED：任务取消成功，UNKNOWN：任务不存在或状态未知", + ) + request_id: Optional[str] = Field( + default=None, + title="Request ID", + description="本次请求的唯一ID，可用于日志追踪。", + ) + + +class ImageToVideoByFirstAndLastFrameWan22Submit( + Component[ + ImageToVideoByFirstAndLastFrameWan22SubmitInput, + ImageToVideoByFirstAndLastFrameWan22SubmitOutput, + ], +): + """ + Submit a keyframe-to-video generation + task using the wan2.2-kf2v-flash model. + """ + + name: str = ( + "modelstudio_image_to_video_by_first_and_last_frame_wan22_submit_task" + ) + description: str = ( + "[版本: wan2.2] 通义万相首尾帧生视频模型（wan2.2-kf2v-flash）异步任务提交工具。\n" + "基于首帧与尾帧图像及文本提示，生成一段流畅的无声视频（当前不支持音频输出）。\n" + ) + + @trace( + trace_type="AIGC", + trace_name="image_to_video_by_first_and_last_frame_wan22_submit", + ) + async def arun( + self, + args: ImageToVideoByFirstAndLastFrameWan22SubmitInput, + **kwargs: Any, + ) -> ImageToVideoByFirstAndLastFrameWan22SubmitOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") + + model_name = kwargs.get( + "model_name", + os.getenv("IMAGE_TO_VIDEO_KF2V_MODEL_NAME", "wan2.2-kf2v-flash"), + ) + + # 构建 parameters（全部为可选参数） + parameters = {} + if args.resolution: + parameters["resolution"] = args.resolution + if args.prompt_extend is not None: + parameters["prompt_extend"] = args.prompt_extend + if args.watermark is not None: + parameters["watermark"] = args.watermark + if args.seed is not None: + parameters["seed"] = args.seed + if args.template: + parameters["template"] = args.template + aio_video_synthesis = AioVideoSynthesis() + + response = await aio_video_synthesis.async_call( + model=model_name, + api_key=api_key, + first_frame_url=args.first_frame_url, + last_frame_url=args.last_frame_url, + prompt=args.prompt, + negative_prompt=args.negative_prompt, + **parameters, + ) + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "results", + "payload": { + "request_id": request_id, + "submit_task": response, + }, + }, + ) + + if ( + response.status_code != HTTPStatus.OK + or not response.output + or response.output.task_status in ["FAILED", "CANCELED"] + ): + raise RuntimeError( + f"Failed to submit keyframe-to-video task: {response}", + ) + + if not request_id: + request_id = ( + response.request_id + if response.request_id + else str(uuid.uuid4()) + ) + + result = ImageToVideoByFirstAndLastFrameWan22SubmitOutput( + request_id=request_id, + task_id=response.output.task_id, + task_status=response.output.task_status, + ) + return result + + +# ========== Fetch 部分 ========== + + +class ImageToVideoByFirstAndLastFrameWan22FetchInput(BaseModel): + task_id: str = Field( + title="Task ID", + description="要查询的视频生成任务ID。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context containing " + "headers for mcp only, don't generate it", + ) + + +class ImageToVideoByFirstAndLastFrameWan22FetchOutput(BaseModel): + video_url: str = Field( + title="Video URL", + description="生成视频的公网可访问URL（MP4格式，无声）。有效期24小时，请及时下载。", + ) + task_id: str = Field( + title="Task ID", + description="任务ID，与输入一致。", + ) + task_status: str = Field( + title="Task Status", + description="任务最终状态，成功时为 SUCCEEDED。", + ) + request_id: Optional[str] = Field( + default=None, + title="Request ID", + description="请求ID，用于追踪。", + ) + + +class ImageToVideoByFirstAndLastFrameWan22Fetch( + Component[ + ImageToVideoByFirstAndLastFrameWan22FetchInput, + ImageToVideoByFirstAndLastFrameWan22FetchOutput, + ], +): + name: str = ( + "modelstudio_image_to_video_by_first_and_last_frame_wan22_fetch_result" + ) + description: str = ( + "查询通义万相 wan2.2-kf2v-flash 首尾帧生视频任务的结果。\n" + "输入 Task ID，返回生成的视频 URL 及任务状态。\n" + "请在提交任务后轮询此接口，直到任务状态变为 SUCCEEDED。\n" + "注意：video_url 有效期为 24 小时。" + ) + + @trace( + trace_type="AIGC", + trace_name="image_to_video_by_first_and_last_frame_wan22_fetch", + ) + async def arun( + self, + args: ImageToVideoByFirstAndLastFrameWan22FetchInput, + **kwargs: Any, + ) -> ImageToVideoByFirstAndLastFrameWan22FetchOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError as e: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") from e + + aio_video_synthesis = AioVideoSynthesis() + + response = await aio_video_synthesis.fetch( + api_key=api_key, + task=args.task_id, + ) + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "results", + "payload": { + "request_id": response.request_id, + "fetch_result": response, + }, + }, + ) + + if ( + response.status_code != HTTPStatus.OK + or not response.output + or response.output.task_status in ["FAILED", "CANCELED"] + ): + raise RuntimeError( + f"Failed to fetch keyframe-to-video result: {response}", + ) + + request_id = response.request_id or request_id or str(uuid.uuid4()) + + return ImageToVideoByFirstAndLastFrameWan22FetchOutput( + video_url=response.output.video_url, + task_id=response.output.task_id, + task_status=response.output.task_status, + request_id=request_id, + ) diff --git a/src/agentscope_bricks/components/generations/async_image_to_video_wan26.py b/src/agentscope_bricks/components/generations/async_image_to_video_wan26.py index 2de8324..3328866 100644 --- a/src/agentscope_bricks/components/generations/async_image_to_video_wan26.py +++ b/src/agentscope_bricks/components/generations/async_image_to_video_wan26.py @@ -1,294 +1,289 @@ -# -*- coding: utf-8 -*- -import os -import uuid -from http import HTTPStatus -from typing import Any, Optional - -from dashscope.aigc.video_synthesis import AioVideoSynthesis -from mcp.server.fastmcp import Context -from pydantic import BaseModel, Field - -from agentscope_bricks.base.component import Component -from agentscope_bricks.utils.tracing_utils.wrapper import trace -from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key -from agentscope_bricks.utils.tracing_utils import TracingUtil - - -class ImageToVideoWan26SubmitInput(BaseModel): - """ - Input model for submitting an image-to-video task using wan2.6-i2v. - """ - - image_url: str = Field( - ..., - description="输入图像，支持公网URL、Base64编码或本地文件路径", - ) - prompt: Optional[str] = Field( - default=None, - description="正向提示词，描述希望视频中发生的动作或变化，例如“镜头缓慢推进，风吹动树叶”。", - ) - negative_prompt: Optional[str] = Field( - default=None, - description="反向提示词，用于排除不希望出现的内容，例如“模糊、闪烁、变形、水印”。", - ) - audio_url: Optional[str] = Field( - default=None, - description="自定义音频文件的公网URL。参数优先级：audio_url > audio。", - ) - audio: Optional[bool] = Field( - default=None, - description="是否自动生成配音。仅在 audio_url 未提供时生效。", - ) - template: Optional[str] = Field( - default=None, - description="视频特效模板，如：squish（解压捏捏）、flying（魔法悬浮）、carousel（时光木马）等。", - ) - resolution: Optional[str] = Field( - default=None, - description="视频分辨率，可选值：'720P'、'1080P'。默认为 '1080P'。", - ) - duration: Optional[int] = Field( - default=None, - description="视频时长（秒），可选值：5、10、15。默认为 5。", - ) - prompt_extend: Optional[bool] = Field( - default=None, - description=" Prompt 智能改写。开启后可提升生成效果，并使 shot_type 生效," - "默认值为 true:开启智能改写。false：不开启智能改写。", - ) - shot_type: Optional[str] = Field( - default=None, - description="镜头类型，仅在 prompt_extend=true 时生效。" - "可选值：'single'（单镜头，默认）、'multi'（多镜头切换）。" - "参数优先级高于 prompt 中的描述。", - ) - watermark: Optional[bool] = Field( - default=None, - description="是否在视频中添加水印（如“AI生成”标识）。默认不添加。", - ) - seed: Optional[int] = Field( - default=None, - description="随机种子，用于结果复现。", - ) - ctx: Optional[Context] = Field( - default=None, - description="HTTP request context containing headers for mcp only, " - "don't generate it", - ) - - -class ImageToVideoWan26SubmitOutput(BaseModel): - """ - Output of the image-to-video task submission. - """ - - task_id: str = Field( - title="Task ID", - description="异步任务的唯一标识符。", - ) - task_status: str = Field( - title="Task Status", - description="视频生成的任务状态，PENDING：任务排队中，RUNNING：任务处理中，SUCCEEDED：任务执行成功，" - "FAILED：任务执行失败，CANCELED：任务取消成功，UNKNOWN：任务不存在或状态未知", - ) - request_id: Optional[str] = Field( - default=None, - title="Request ID", - description="本次请求的唯一ID，可用于日志追踪。", - ) - - -class ImageToVideoWan26Submit( - Component[ImageToVideoWan26SubmitInput, ImageToVideoWan26SubmitOutput], -): - """ - Submit an image-to-video generation task using the wan2.6-i2v model. - """ - - name: str = "modelstudio_image_to_video_wan26_submit_task" - description: str = ( - "[版本: wan2.6] 通义万相图生视频模型（wan2.6-i2v）异步任务提交工具。基于单张首帧图像和文本提示，生成一段流畅的有声视频。\n" # noqa - "支持视频时长：5秒、10秒或15秒；分辨率：720P、1080P；支持自动配音或传入自定义音频，实现音画同步。\n" - "独家支持多镜头叙事：可生成包含多个镜头的视频，并在镜头切换时保持主体一致性。\n" - "提供特效模板（如“魔法悬浮”、“气球膨胀”），适用于创意视频制作、娱乐特效展示等场景。\n" - ) - - @trace(trace_type="AIGC", trace_name="image_to_video_wan26_submit") - async def arun( - self, - args: ImageToVideoWan26SubmitInput, - **kwargs: Any, - ) -> ImageToVideoWan26SubmitOutput: - trace_event = kwargs.pop("trace_event", None) - request_id = TracingUtil.get_request_id() - - try: - api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) - except AssertionError: - raise ValueError("Please set valid DASHSCOPE_API_KEY!") - - model_name = kwargs.get( - "model_name", - os.getenv("IMAGE_TO_VIDEO_MODEL_NAME", "wan2.6-i2v"), - ) - - # 构建 parameters（全部为可选参数） - parameters = {} - if args.audio is not None: - parameters["audio"] = args.audio - if args.resolution: - parameters["resolution"] = args.resolution - if args.duration is not None: - parameters["duration"] = args.duration - if args.prompt_extend is not None: - parameters["prompt_extend"] = args.prompt_extend - if args.shot_type: - parameters["shot_type"] = args.shot_type - if args.watermark is not None: - parameters["watermark"] = args.watermark - if args.seed is not None: - parameters["seed"] = args.seed - aio_video_synthesis = AioVideoSynthesis() - - response = await aio_video_synthesis.async_call( - model=model_name, - api_key=api_key, - img_url=args.image_url, - prompt=args.prompt, - negative_prompt=args.negative_prompt, - audio_url=args.audio_url, - template=args.template, - **parameters, - ) - - if trace_event: - trace_event.on_log( - "", - **{ - "step_suffix": "results", - "payload": { - "request_id": request_id, - "submit_task": response, - }, - }, - ) - - if ( - response.status_code != HTTPStatus.OK - or not response.output - or response.output.task_status in ["FAILED", "CANCELED"] - ): - raise RuntimeError( - f"Failed to submit image-to-video task: {response}", - ) - - if not request_id: - request_id = ( - response.request_id - if response.request_id - else str(uuid.uuid4()) - ) - - result = ImageToVideoWan26SubmitOutput( - request_id=request_id, - task_id=response.output.task_id, - task_status=response.output.task_status, - ) - return result - - -# ========== Fetch 部分保持不变（仅微调描述） ========== - - -class ImageToVideoWan26FetchInput(BaseModel): # noqa - task_id: str = Field( - title="Task ID", - description="要查询的视频生成任务ID。", - ) - ctx: Optional[Context] = Field( - default=None, - description="HTTP request context containing headers for mcp only, " - "don't generate it", - ) - - -class ImageToVideoWan26FetchOutput(BaseModel): - video_url: str = Field( - title="Video URL", - description="生成视频的公网可访问URL（MP4格式）。", - ) - task_id: str = Field( - title="Task ID", - description="任务ID，与输入一致。", - ) - task_status: str = Field( - title="Task Status", - description="任务最终状态，成功时为 SUCCEEDED。", - ) - request_id: Optional[str] = Field( - default=None, - title="Request ID", - description="请求ID，用于追踪。", - ) - - -class ImageToVideoWan26Fetch( - Component[ImageToVideoWan26FetchInput, ImageToVideoWan26FetchOutput], -): - name: str = "modelstudio_image_to_video_wan26_fetch_result" - description: str = ( - "查询通义万相 wan2.6-i2v 图生视频任务的结果。" - "输入 Task ID，返回生成的视频 URL 及任务状态。" - "请在提交任务后轮询此接口，直到任务状态变为 SUCCEEDED。" - ) - - @trace(trace_type="AIGC", trace_name="image_to_video_wan26_fetch") - async def arun( - self, - args: ImageToVideoWan26FetchInput, - **kwargs: Any, - ) -> ImageToVideoWan26FetchOutput: - trace_event = kwargs.pop("trace_event", None) - request_id = TracingUtil.get_request_id() - - try: - api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) - except AssertionError as e: - raise ValueError("Please set valid DASHSCOPE_API_KEY!") from e - - aio_video_synthesis = AioVideoSynthesis() - - response = await aio_video_synthesis.fetch( - api_key=api_key, - task=args.task_id, - ) - - if trace_event: - trace_event.on_log( - "", - **{ - "step_suffix": "results", - "payload": { - "request_id": response.request_id, - "fetch_result": response, - }, - }, - ) - - if ( - response.status_code != HTTPStatus.OK - or not response.output - or response.output.task_status in ["FAILED", "CANCELED"] - ): - raise RuntimeError( - f"Failed to fetch image-to-video result: {response}", - ) - - request_id = response.request_id or request_id or str(uuid.uuid4()) - - return ImageToVideoWan26FetchOutput( - video_url=response.output.video_url, - task_id=response.output.task_id, - task_status=response.output.task_status, - request_id=request_id, - ) +# -*- coding: utf-8 -*- +import os +import uuid +from http import HTTPStatus +from typing import Any, Optional + +from dashscope.aigc.video_synthesis import AioVideoSynthesis +from mcp.server.fastmcp import Context +from pydantic import BaseModel, Field + +from agentscope_bricks.base.component import Component +from agentscope_bricks.utils.tracing_utils.wrapper import trace +from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key +from agentscope_bricks.utils.tracing_utils import TracingUtil + + +class ImageToVideoWan26SubmitInput(BaseModel): + """ + Input model for submitting an image-to-video task using wan2.6-i2v. + """ + + image_url: str = Field( + ..., + description="输入图像，支持公网URL、Base64编码", + ) + prompt: Optional[str] = Field( + default=None, + description="正向提示词，描述希望视频中发生的动作或变化，例如“镜头缓慢推进，风吹动树叶”。", + ) + negative_prompt: Optional[str] = Field( + default=None, + description="反向提示词，用于排除不希望出现的内容，例如“模糊、闪烁、变形、水印”。", + ) + audio_url: Optional[str] = Field( + default=None, + description="自定义音频文件的公网URL。参数优先级：audio_url > audio。" + "若不提供audio_url ，模型将根据视频内容自动生成匹配的背景音乐或音效。", + ) + template: Optional[str] = Field( + default=None, + description="视频特效模板，如：flying，表示使用“魔法悬浮”特效等。", + ) + resolution: Optional[str] = Field( + default=None, + description="视频分辨率，可选值：'720P'、'1080P'。默认为 '1080P'。", + ) + duration: Optional[int] = Field( + default=None, + description="视频时长（秒），可选值：5、10、15。默认为 5。", + ) + prompt_extend: Optional[bool] = Field( + default=None, + description=" Prompt 智能改写。开启后可提升生成效果，并使 shot_type 生效," + "默认值为 true:开启智能改写。false：不开启智能改写。", + ) + shot_type: Optional[str] = Field( + default=None, + description="镜头类型，仅在 prompt_extend=true 时生效。" + "可选值：'single'（单镜头，默认）、'multi'（多镜头切换）。" + "参数优先级高于 prompt 中的描述。", + ) + watermark: Optional[bool] = Field( + default=None, + description="是否添加水印,false：默认值，不添加水印,true：添加水印。", + ) + seed: Optional[int] = Field( + default=None, + description="随机种子，用于结果复现。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context containing headers for mcp only, " + "don't generate it", + ) + + +class ImageToVideoWan26SubmitOutput(BaseModel): + """ + Output of the image-to-video task submission. + """ + + task_id: str = Field( + title="Task ID", + description="异步任务的唯一标识符。", + ) + task_status: str = Field( + title="Task Status", + description="视频生成的任务状态，PENDING：任务排队中，RUNNING：任务处理中，SUCCEEDED：任务执行成功，" + "FAILED：任务执行失败，CANCELED：任务取消成功，UNKNOWN：任务不存在或状态未知", + ) + request_id: Optional[str] = Field( + default=None, + title="Request ID", + description="本次请求的唯一ID，可用于日志追踪。", + ) + + +class ImageToVideoWan26Submit( + Component[ImageToVideoWan26SubmitInput, ImageToVideoWan26SubmitOutput], +): + """ + Submit an image-to-video generation task using the wan2.6-i2v model. + """ + + name: str = "modelstudio_image_to_video_wan26_submit_task" + description: str = ( + "[版本: wan2.6] 通义万相图生视频模型（wan2.6-i2v）异步任务提交工具。基于单张首帧图像和文本提示，生成一段流畅的有声视频。\n" # noqa + "支持视频时长：5秒、10秒或15秒；分辨率：720P、1080P；支持自动配音或传入自定义音频，实现音画同步。\n" + "独家支持多镜头叙事：可生成包含多个镜头的视频，并在镜头切换时保持主体一致性。\n" + "提供特效模板（如“魔法悬浮”、“气球膨胀”），适用于创意视频制作、娱乐特效展示等场景。\n" + ) + + @trace(trace_type="AIGC", trace_name="image_to_video_wan26_submit") + async def arun( + self, + args: ImageToVideoWan26SubmitInput, + **kwargs: Any, + ) -> ImageToVideoWan26SubmitOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") + + model_name = kwargs.get( + "model_name", + os.getenv("IMAGE_TO_VIDEO_MODEL_NAME", "wan2.6-i2v"), + ) + + # 构建 parameters（全部为可选参数） + parameters = {} + if args.resolution: + parameters["resolution"] = args.resolution + if args.duration is not None: + parameters["duration"] = args.duration + if args.prompt_extend is not None: + parameters["prompt_extend"] = args.prompt_extend + if args.shot_type: + parameters["shot_type"] = args.shot_type + if args.watermark is not None: + parameters["watermark"] = args.watermark + if args.seed is not None: + parameters["seed"] = args.seed + aio_video_synthesis = AioVideoSynthesis() + + response = await aio_video_synthesis.async_call( + model=model_name, + api_key=api_key, + img_url=args.image_url, + prompt=args.prompt, + negative_prompt=args.negative_prompt, + audio_url=args.audio_url, + template=args.template, + **parameters, + ) + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "results", + "payload": { + "request_id": request_id, + "submit_task": response, + }, + }, + ) + + if ( + response.status_code != HTTPStatus.OK + or not response.output + or response.output.task_status in ["FAILED", "CANCELED"] + ): + raise RuntimeError( + f"Failed to submit image-to-video task: {response}", + ) + + if not request_id: + request_id = ( + response.request_id + if response.request_id + else str(uuid.uuid4()) + ) + + result = ImageToVideoWan26SubmitOutput( + request_id=request_id, + task_id=response.output.task_id, + task_status=response.output.task_status, + ) + return result + + +# ========== Fetch 部分保持不变（仅微调描述） ========== + + +class ImageToVideoWan26FetchInput(BaseModel): # noqa + task_id: str = Field( + title="Task ID", + description="要查询的视频生成任务ID。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context containing headers for mcp only, " + "don't generate it", + ) + + +class ImageToVideoWan26FetchOutput(BaseModel): + video_url: str = Field( + title="Video URL", + description="生成视频的公网可访问URL（MP4格式）。", + ) + task_id: str = Field( + title="Task ID", + description="任务ID，与输入一致。", + ) + task_status: str = Field( + title="Task Status", + description="任务最终状态，成功时为 SUCCEEDED。", + ) + request_id: Optional[str] = Field( + default=None, + title="Request ID", + description="请求ID，用于追踪。", + ) + + +class ImageToVideoWan26Fetch( + Component[ImageToVideoWan26FetchInput, ImageToVideoWan26FetchOutput], +): + name: str = "modelstudio_image_to_video_wan26_fetch_result" + description: str = ( + "查询通义万相 wan2.6-i2v 图生视频任务的结果。" + "输入 Task ID，返回生成的视频 URL 及任务状态。" + "请在提交任务后轮询此接口，直到任务状态变为 SUCCEEDED。" + ) + + @trace(trace_type="AIGC", trace_name="image_to_video_wan26_fetch") + async def arun( + self, + args: ImageToVideoWan26FetchInput, + **kwargs: Any, + ) -> ImageToVideoWan26FetchOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError as e: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") from e + + aio_video_synthesis = AioVideoSynthesis() + + response = await aio_video_synthesis.fetch( + api_key=api_key, + task=args.task_id, + ) + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "results", + "payload": { + "request_id": response.request_id, + "fetch_result": response, + }, + }, + ) + + if ( + response.status_code != HTTPStatus.OK + or not response.output + or response.output.task_status in ["FAILED", "CANCELED"] + ): + raise RuntimeError( + f"Failed to fetch image-to-video result: {response}", + ) + + request_id = response.request_id or request_id or str(uuid.uuid4()) + + return ImageToVideoWan26FetchOutput( + video_url=response.output.video_url, + task_id=response.output.task_id, + task_status=response.output.task_status, + request_id=request_id, + ) diff --git a/src/agentscope_bricks/components/generations/async_text_to_video_wan26.py b/src/agentscope_bricks/components/generations/async_text_to_video_wan26.py index a446597..f14faa9 100644 --- a/src/agentscope_bricks/components/generations/async_text_to_video_wan26.py +++ b/src/agentscope_bricks/components/generations/async_text_to_video_wan26.py @@ -1,298 +1,291 @@ -# -*- coding: utf-8 -*- -import os -import uuid -from http import HTTPStatus -from typing import Any, Optional - -from dashscope.aigc.video_synthesis import AioVideoSynthesis -from mcp.server.fastmcp import Context -from pydantic import BaseModel, Field - -from agentscope_bricks.base.component import Component -from agentscope_bricks.utils.tracing_utils.wrapper import trace -from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key -from agentscope_bricks.utils.tracing_utils import TracingUtil - - -class TextToVideoWan26SubmitInput(BaseModel): - """ - Input model for text-to-video generation submission using wan2.6-t2v. - """ - - prompt: str = Field( - ..., - description="正向提示词，描述希望生成的视频内容，例如“一只宇航员猫在火星上跳舞”", - ) - negative_prompt: Optional[str] = Field( - default=None, - description="反向提示词，描述不希望出现在视频中的内容，例如“模糊、水印、文字、变形”", - ) - audio_url: Optional[str] = Field( - default=None, - description="自定义音频文件URL，模型将使用该音频生成视频。" - "参数优先级：audio_url > audio，仅在 audio_url 为空时 audio 生效。", - ) - audio: Optional[bool] = Field( - default=None, - description="是否自动生成音频。" - "参数优先级：audio_url > audio，仅在 audio_url 为空时 audio 生效。", - ) - size: Optional[str] = Field( - default=None, - description="视频分辨率，默认值为 1920*1080（具体支持值请参考文档）", - ) - duration: Optional[int] = Field( - default=None, - description="视频时长（秒），可选值：5、10、15。默认为 5。", - ) - prompt_extend: Optional[bool] = Field( - default=None, - description="是否开启prompt智能改写，开启后使用大模型对输入prompt进行智能优化", - ) - shot_type: Optional[str] = Field( - default=None, - description="镜头类型，仅在 prompt_extend=true 时生效。" - "可选值：'single'（单镜头，默认）、'multi'（多镜头切换）。" - "参数优先级高于 prompt 中的描述。", - ) - watermark: Optional[bool] = Field( - default=None, - description="是否添加水印，默认不设置", - ) - seed: Optional[int] = Field( - default=None, - description="随机种子，用于结果复现。", - ) - ctx: Optional[Context] = Field( - default=None, - description="HTTP request context containing headers " - "for mcp only, don't generate it", - ) - - -class TextToVideoWan26SubmitOutput(BaseModel): - """ - Output model for text-to-video generation submission. - """ - - task_id: str = Field( - title="Task ID", - description="视频生成的任务ID", - ) - - task_status: str = Field( - title="Task Status", - description="任务状态：PENDING（排队中）、RUNNING（处理中）、SUCCEEDED（成功）、" - "FAILED（失败）、CANCELED（已取消）、UNKNOWN（未知）", - ) - - request_id: Optional[str] = Field( - default=None, - title="Request ID", - description="请求ID，用于追踪", - ) - - -class TextToVideoWan26Submit( - Component[TextToVideoWan26SubmitInput, TextToVideoWan26SubmitOutput], -): - """ - Service for submitting text-to-video - generation tasks using Wan 2.6 T2V model. - """ - - name: str = "modelstudio_text_to_video_wan26_submit_task" - description: str = ( - "[版本: wan2.6] 通义万相文生视频模型（wan2.6-t2v）异步任务提交工具。基于纯文本提示生成一段流畅的有声视频。\n" - "支持视频时长：5秒、10秒或15秒；分辨率：720P、1080P；支持自动配音或传入自定义音频，实现音画同步。\n" - "独家支持多镜头叙事：可生成包含多个镜头的视频，并在镜头切换时保持主体一致性。\n" - ) - - @trace(trace_type="AIGC", trace_name="text_to_video_wan26_submit") - async def arun( - self, - args: TextToVideoWan26SubmitInput, - **kwargs: Any, - ) -> TextToVideoWan26SubmitOutput: - trace_event = kwargs.pop("trace_event", None) - request_id = TracingUtil.get_request_id() - - try: - api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) - except AssertionError: - raise ValueError("Please set valid DASHSCOPE_API_KEY!") - - model_name = kwargs.get( - "model_name", - os.getenv("TEXT_TO_VIDEO_MODEL_NAME", "wan2.6-t2v"), - ) - - parameters = {} - if args.audio is not None: - parameters["audio"] = args.audio - if args.size: - parameters["size"] = args.size - if args.duration is not None: - parameters["duration"] = args.duration - if args.prompt_extend is not None: - parameters["prompt_extend"] = args.prompt_extend - if args.watermark is not None: - parameters["watermark"] = args.watermark - if args.shot_type: - parameters["shot_type"] = args.shot_type - if args.seed is not None: - parameters["seed"] = args.seed - aio_video_synthesis = AioVideoSynthesis() - - response = await aio_video_synthesis.async_call( - model=model_name, - api_key=api_key, - prompt=args.prompt, - negative_prompt=args.negative_prompt, - audio_url=args.audio_url, - **parameters, - ) - - if trace_event: - trace_event.on_log( - "", - **{ - "step_suffix": "results", - "payload": { - "request_id": request_id, - "submit_task": response, - }, - }, - ) - - if ( - response.status_code != HTTPStatus.OK - or not response.output - or response.output.task_status in ["FAILED", "CANCELED"] - ): - raise RuntimeError( - f"Failed to submit text-to-video task: {response}", - ) - - if not request_id: - request_id = response.request_id or str(uuid.uuid4()) - - result = TextToVideoWan26SubmitOutput( - request_id=request_id, - task_id=response.output.task_id, - task_status=response.output.task_status, - ) - return result - - -class TextToVideoWan26FetchInput(BaseModel): - """ - Input model for fetching text-to-video generation results. - """ - - task_id: str = Field( - title="Task ID", - description="视频生成的任务ID", - ) - ctx: Optional[Context] = Field( - default=None, - description="HTTP request context containing headers " - "for mcp only, don't generate it", - ) - - -class TextToVideoWan26FetchOutput(BaseModel): - """ - Output model for fetching text-to-video generation results. - """ - - video_url: str = Field( - title="Video URL", - description="生成的视频公网可访问URL", - ) - - task_id: str = Field( - title="Task ID", - description="视频生成的任务ID", - ) - - task_status: str = Field( - title="Task Status", - description="任务状态：PENDING、RUNNING、SUCCEEDED、FAILED、CANCELED、UNKNOWN", - ) - - request_id: Optional[str] = Field( - default=None, - title="Request ID", - description="请求ID", - ) - - -class TextToVideoWan26Fetch( - Component[TextToVideoWan26FetchInput, TextToVideoWan26FetchOutput], -): - """ - Service for fetching text-to-video generation results. - """ - - name: str = "modelstudio_text_to_video_wan26_fetch_result" - description: str = ( - "通义万相-文生视频模型（wan2.6-t2v）的异步任务结果查询工具，根据Task ID查询生成的视频URL。" - ) - - @trace(trace_type="AIGC", trace_name="text_to_video_wan26_fetch") - async def arun( - self, - args: TextToVideoWan26FetchInput, - **kwargs: Any, - ) -> TextToVideoWan26FetchOutput: - trace_event = kwargs.pop("trace_event", None) - request_id = TracingUtil.get_request_id() - - try: - api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) - except AssertionError as e: - raise ValueError("Please set valid DASHSCOPE_API_KEY!") from e - - aio_video_synthesis = AioVideoSynthesis() - - response = await aio_video_synthesis.fetch( - api_key=api_key, - task=args.task_id, - ) - - if trace_event: - trace_event.on_log( - "", - **{ - "step_suffix": "results", - "payload": { - "request_id": response.request_id, - "fetch_result": response, - }, - }, - ) - - if ( - response.status_code != HTTPStatus.OK - or not response.output - or response.output.task_status in ["FAILED", "CANCELED"] - ): - raise RuntimeError( - f"Failed to fetch text-to-video result: {response}", - ) - - if not request_id: - request_id = ( - response.request_id - if response.request_id - else str(uuid.uuid4()) - ) - - result = TextToVideoWan26FetchOutput( - video_url=response.output.video_url, - task_id=response.output.task_id, - task_status=response.output.task_status, - request_id=request_id, - ) - return result +# -*- coding: utf-8 -*- +import os +import uuid +from http import HTTPStatus +from typing import Any, Optional + +from dashscope.aigc.video_synthesis import AioVideoSynthesis +from mcp.server.fastmcp import Context +from pydantic import BaseModel, Field + +from agentscope_bricks.base.component import Component +from agentscope_bricks.utils.tracing_utils.wrapper import trace +from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key +from agentscope_bricks.utils.tracing_utils import TracingUtil + + +class TextToVideoWan26SubmitInput(BaseModel): + """ + Input model for text-to-video generation submission using wan2.6-t2v. + """ + + prompt: str = Field( + ..., + description="正向提示词，描述希望生成的视频内容，例如“一只宇航员猫在火星上跳舞”", + ) + negative_prompt: Optional[str] = Field( + default=None, + description="反向提示词，描述不希望出现在视频中的内容，例如“模糊、水印、文字、变形”", + ) + audio_url: Optional[str] = Field( + default=None, + description="自定义音频文件URL，模型将使用该音频生成视频。" + "若不提供audio_url ，模型将根据视频内容自动生成匹配的背景音乐或音效。", + ) + size: Optional[str] = Field( + default=None, + description="视频分辨率，默认值为 1920*1080（具体支持值请参考文档）", + ) + duration: Optional[int] = Field( + default=None, + description="视频时长（秒），可选值：5、10、15。默认为 5。", + ) + prompt_extend: Optional[bool] = Field( + default=None, + description="是否开启prompt智能改写，开启后使用大模型对输入prompt进行智能优化", + ) + shot_type: Optional[str] = Field( + default=None, + description="镜头类型，仅在 prompt_extend=true 时生效。" + "可选值：'single'（单镜头，默认）、'multi'（多镜头切换）。" + "参数优先级高于 prompt 中的描述。", + ) + watermark: Optional[bool] = Field( + default=None, + description="是否在视频中添加水印,false：默认值，不添加水印,true：添加水印。", + ) + seed: Optional[int] = Field( + default=None, + description="随机种子，用于结果复现。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context containing headers " + "for mcp only, don't generate it", + ) + + +class TextToVideoWan26SubmitOutput(BaseModel): + """ + Output model for text-to-video generation submission. + """ + + task_id: str = Field( + title="Task ID", + description="视频生成的任务ID", + ) + + task_status: str = Field( + title="Task Status", + description="任务状态：PENDING（排队中）、RUNNING（处理中）、SUCCEEDED（成功）、" + "FAILED（失败）、CANCELED（已取消）、UNKNOWN（未知）", + ) + + request_id: Optional[str] = Field( + default=None, + title="Request ID", + description="请求ID，用于追踪", + ) + + +class TextToVideoWan26Submit( + Component[TextToVideoWan26SubmitInput, TextToVideoWan26SubmitOutput], +): + """ + Service for submitting text-to-video + generation tasks using Wan 2.6 T2V model. + """ + + name: str = "modelstudio_text_to_video_wan26_submit_task" + description: str = ( + "[版本: wan2.6] 通义万相文生视频模型（wan2.6-t2v）异步任务提交工具。基于纯文本提示生成一段流畅的有声视频。\n" + "支持视频时长：5秒、10秒或15秒；分辨率：720P、1080P；支持自动配音或传入自定义音频，实现音画同步。\n" + "独家支持多镜头叙事：可生成包含多个镜头的视频，并在镜头切换时保持主体一致性。\n" + ) + + @trace(trace_type="AIGC", trace_name="text_to_video_wan26_submit") + async def arun( + self, + args: TextToVideoWan26SubmitInput, + **kwargs: Any, + ) -> TextToVideoWan26SubmitOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") + + model_name = kwargs.get( + "model_name", + os.getenv("TEXT_TO_VIDEO_MODEL_NAME", "wan2.6-t2v"), + ) + + parameters = {} + if args.size: + parameters["size"] = args.size + if args.duration is not None: + parameters["duration"] = args.duration + if args.prompt_extend is not None: + parameters["prompt_extend"] = args.prompt_extend + if args.watermark is not None: + parameters["watermark"] = args.watermark + if args.shot_type: + parameters["shot_type"] = args.shot_type + if args.seed is not None: + parameters["seed"] = args.seed + aio_video_synthesis = AioVideoSynthesis() + + response = await aio_video_synthesis.async_call( + model=model_name, + api_key=api_key, + prompt=args.prompt, + negative_prompt=args.negative_prompt, + audio_url=args.audio_url, + **parameters, + ) + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "results", + "payload": { + "request_id": request_id, + "submit_task": response, + }, + }, + ) + + if ( + response.status_code != HTTPStatus.OK + or not response.output + or response.output.task_status in ["FAILED", "CANCELED"] + ): + raise RuntimeError( + f"Failed to submit text-to-video task: {response}", + ) + + if not request_id: + request_id = response.request_id or str(uuid.uuid4()) + + result = TextToVideoWan26SubmitOutput( + request_id=request_id, + task_id=response.output.task_id, + task_status=response.output.task_status, + ) + return result + + +class TextToVideoWan26FetchInput(BaseModel): + """ + Input model for fetching text-to-video generation results. + """ + + task_id: str = Field( + title="Task ID", + description="视频生成的任务ID", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context containing headers " + "for mcp only, don't generate it", + ) + + +class TextToVideoWan26FetchOutput(BaseModel): + """ + Output model for fetching text-to-video generation results. + """ + + video_url: str = Field( + title="Video URL", + description="生成的视频公网可访问URL", + ) + + task_id: str = Field( + title="Task ID", + description="视频生成的任务ID", + ) + + task_status: str = Field( + title="Task Status", + description="任务状态：PENDING、RUNNING、SUCCEEDED、FAILED、CANCELED、UNKNOWN", + ) + + request_id: Optional[str] = Field( + default=None, + title="Request ID", + description="请求ID", + ) + + +class TextToVideoWan26Fetch( + Component[TextToVideoWan26FetchInput, TextToVideoWan26FetchOutput], +): + """ + Service for fetching text-to-video generation results. + """ + + name: str = "modelstudio_text_to_video_wan26_fetch_result" + description: str = ( + "通义万相-文生视频模型（wan2.6-t2v）的异步任务结果查询工具，根据Task ID查询生成的视频URL。" + ) + + @trace(trace_type="AIGC", trace_name="text_to_video_wan26_fetch") + async def arun( + self, + args: TextToVideoWan26FetchInput, + **kwargs: Any, + ) -> TextToVideoWan26FetchOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError as e: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") from e + + aio_video_synthesis = AioVideoSynthesis() + + response = await aio_video_synthesis.fetch( + api_key=api_key, + task=args.task_id, + ) + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "results", + "payload": { + "request_id": response.request_id, + "fetch_result": response, + }, + }, + ) + + if ( + response.status_code != HTTPStatus.OK + or not response.output + or response.output.task_status in ["FAILED", "CANCELED"] + ): + raise RuntimeError( + f"Failed to fetch text-to-video result: {response}", + ) + + if not request_id: + request_id = ( + response.request_id + if response.request_id + else str(uuid.uuid4()) + ) + + result = TextToVideoWan26FetchOutput( + video_url=response.output.video_url, + task_id=response.output.task_id, + task_status=response.output.task_status, + request_id=request_id, + ) + return result diff --git a/src/agentscope_bricks/components/generations/image_edit_wan26.py b/src/agentscope_bricks/components/generations/image_edit_wan26.py new file mode 100644 index 0000000..855aa2f --- /dev/null +++ b/src/agentscope_bricks/components/generations/image_edit_wan26.py @@ -0,0 +1,211 @@ +# -*- coding: utf-8 -*- +import uuid +from typing import Any, Optional +from dashscope import AioMultiModalConversation +from mcp.server.fastmcp import Context +from pydantic import BaseModel, Field + +from agentscope_bricks.base.component import Component +from agentscope_bricks.utils.tracing_utils.wrapper import trace, TraceType +from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key +from agentscope_bricks.utils.tracing_utils import TracingUtil + + +class ImageGenInput(BaseModel): + """ + Input schema for Wanx 2.6 image editing generation. + """ + + prompt: str = Field( + ..., + description="正向提示词，描述期望生成的图像内容", + ) + negative_prompt: Optional[str] = Field( + default=None, + description="反向提示词，描述不希望出现的内容，如低质量、模糊、文字等。", + ) + size: Optional[str] = Field( + default=None, + description="输出图像的分辨率。默认值是1280*1280，可不填。", + ) + prompt_extend: Optional[bool] = Field( + default=None, + description="是否开启 Prompt 智能改写。将使用大模型优化正向提示词。true: 开启（默认），false：不开启。", + ) + seed: Optional[int] = Field( + default=None, + description="随机种子，用于结果复现。", + ) + watermark: Optional[bool] = Field( + default=None, + description="是否添加水印,false：默认值，不添加水印,true：添加水印。", + ) + n: Optional[int] = Field( + default=4, + description="生成图片的数量。取值范围为1~4张 默认4", + ) + images: list[str] = Field( + ..., + description=( + "参考图像URL列表，用于图像编辑。\n" "必须提供至少1张参考图像。" + ), + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context for " + "MCP internal use only, do not generate it.", + ) + + +class ImageGenOutput(BaseModel): + """ + Output schema for Wanx 2.6 image generation. + """ + + results: list[str] = Field( + title="Results", + description="生成的图片URL列表。", + ) + request_id: Optional[str] = Field( + default=None, + title="Request ID", + description="本次请求的唯一标识。", + ) + + +class ImageEditWan26( + Component[ImageGenInput, ImageGenOutput], +): + """ + Wanx 2.6 Image Editing Generation Tool. + Supports: + - Image editing mode (with 1-3 reference images) + Uses the 'wan2.6-image' model from DashScope. + """ + + name: str = "modelstudio_image_edit_wan26" + description: str = ( + "[版本: wan2.6] 通义万相文生图模型（wan2.6-image）。\n" + "图像编辑,基于1～4张输入图像进行编辑、风格迁移或主体一致性生成。返回编辑后的图片URL列表。" + ) + + @trace(trace_type=TraceType.AIGC, trace_name="wanx26_image_generation") + async def arun( + self, + args: ImageGenInput, + **kwargs: Any, + ) -> ImageGenOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") + + model_name = "wan2.6-image" + + # 构造多模态 content：文本 + 可选图像 + content = [{"text": args.prompt}] + images = args.images or [] # 安全处理 None + for img_url in images: + content.append({"image": img_url}) + + messages = [ + { + "role": "user", + "content": content, + }, + ] + parameters = {} + if args.negative_prompt: + parameters["negative_prompt"] = args.negative_prompt + if args.size and args.size != "1280*1280": + parameters["size"] = args.size + if args.seed is not None: + parameters["seed"] = args.seed + if args.watermark is not None: + parameters["watermark"] = args.watermark + if args.prompt_extend is not None: + parameters["prompt_extend"] = args.prompt_extend + if args.n is not None and args.n != 4: + parameters["n"] = args.n + try: + response = await AioMultiModalConversation.call( + api_key=api_key, + model=model_name, + messages=messages, + enable_interleave=False, + **parameters, + ) + except Exception as e: + raise RuntimeError( + f"Failed to call Wanx 2.6 image generation API: {str(e)}", + ) from e + + if response.status_code != 200 or not response.output: + raise RuntimeError(f"Wanx 2.6 image generation failed: {response}") + + results = [] + + try: + if hasattr(response, "output") and response.output: + choices = getattr(response.output, "choices", []) + if choices: + for choice in choices: + message = getattr(choice, "message", {}) + msg_content = getattr(message, "content", []) + if isinstance(msg_content, list): + # 遍历当前 choice 的 content + for item in msg_content: + if isinstance(item, dict) and "image" in item: + results.append(item["image"]) + elif isinstance(item, str) and item.startswith( + ("http://", "https://"), + ): + results.append(item) + elif isinstance( + msg_content, + str, + ) and msg_content.startswith( + ("http://", "https://"), + ): + results.append(msg_content) + elif ( + isinstance(msg_content, dict) + and "image" in msg_content + ): + results.append(msg_content["image"]) + # --- 修改结束 --- + except Exception as e: + raise RuntimeError( + f"Failed to parse Wanx 2.6 API response: {str(e)}", + ) from e + + if not results: + raise RuntimeError(f"No image found in response: {response}") + + if not request_id: + request_id = getattr(response, "request_id", None) or str( + uuid.uuid4(), + ) + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "results", + "payload": { + "request_id": request_id, + "wanx26_image_generation_result": { + "status_code": response.status_code, + "results": results, + }, + }, + }, + ) + + return ImageGenOutput( + results=results, + request_id=request_id, + ) diff --git a/src/agentscope_bricks/components/generations/image_generation_wan26.py b/src/agentscope_bricks/components/generations/image_generation_wan26.py index 21def3c..efe20cb 100644 --- a/src/agentscope_bricks/components/generations/image_generation_wan26.py +++ b/src/agentscope_bricks/components/generations/image_generation_wan26.py @@ -1,194 +1,195 @@ -# -*- coding: utf-8 -*- -import uuid -from typing import Any, Optional -from dashscope import AioMultiModalConversation -from mcp.server.fastmcp import Context -from pydantic import BaseModel, Field - -from agentscope_bricks.base.component import Component -from agentscope_bricks.utils.tracing_utils.wrapper import trace, TraceType -from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key -from agentscope_bricks.utils.tracing_utils import TracingUtil - - -class ImageGenerationWan26Input(BaseModel): - """ - Input schema for Wanx 2.6 text-to-image generation. - """ - - prompt: str = Field( - ..., - description="正向提示词，描述期望生成的图像内容，建议详细且清晰。超过800字符将被截断。", - ) - negative_prompt: Optional[str] = Field( - default=None, - description="反向提示词，描述不希望出现的内容，如低质量、模糊、文字等。超过500字符将被截断。", - ) - size: Optional[str] = Field( - default=None, - description="输出图像的分辨率。默认值是1280*1280，可不填。", - ) - prompt_extend: Optional[bool] = Field( - default=None, - description="是否开启 Prompt 智能改写。将使用大模型优化正向提示词。true: 开启（默认），false：不开启。", - ) - n: Optional[int] = Field( - default=1, - description="生成图片的数量。取值范围为1~4张 默认1", - ) - seed: Optional[int] = Field( - default=None, - description="随机种子，用于结果复现。", - ) - watermark: Optional[bool] = Field( - default=None, - description="是否添加阿里云水印，默认不添加。", - ) - ctx: Optional[Context] = Field( - default=None, - description="HTTP request context for " - "MCP internal use only, do not generate it.", - ) - - -class ImageGenerationWan26Output(BaseModel): - """ - Output schema for Wanx 2.6 text-to-image generation. - """ - - results: list[str] = Field( - title="Results", - description="生成的图片URL列表。", - ) - request_id: Optional[str] = Field( - default=None, - title="Request ID", - description="本次请求的唯一标识。", - ) - - -class ImageGenerationWan26( - Component[ImageGenerationWan26Input, ImageGenerationWan26Output], -): - """ - Wanx 2.6 Text-to-Image Generation Tool. - Uses the 'wan2.6-t2i' model from DashScope - to generate high-quality images from text. - """ - - name: str = "modelstudio_wanx26_image_generation" - description: str = ( - "[版本: wan2.6] 通义万相文生图模型（wanx2.6-t2i）。AI绘画服务，根据文本描述生成高质量图像，并返回图片URL。\n" - "新功能包括图像编辑和图文混合输出，满足更多样化的生成与集成需求。\n" - "支持自定义分辨率：图像面积介于 768×768 至 1440×1440 像素之间，" - "允许在该范围内自由调整宽高比（例如 768×2700）。\n" - ) - - @trace(trace_type=TraceType.AIGC, trace_name="wanx26_image_generation") - async def arun( - self, - args: ImageGenerationWan26Input, - **kwargs: Any, - ) -> ImageGenerationWan26Output: - trace_event = kwargs.pop("trace_event", None) - request_id = TracingUtil.get_request_id() - - try: - api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) - except AssertionError: - raise ValueError("Please set valid DASHSCOPE_API_KEY!") - - model_name = "wan2.6-t2i" - messages = [ - { - "role": "user", - "content": [{"text": args.prompt}], - }, - ] - - # Normalize watermark - if args.watermark is not None: - if isinstance(args.watermark, str): - args.watermark = args.watermark.strip().lower() in ( - "true", - "1", - ) - else: - args.watermark = bool(args.watermark) - - parameters = {} - if args.negative_prompt: - parameters["negative_prompt"] = args.negative_prompt - if args.size and args.size != "1024*1024": - parameters["size"] = args.size - if args.n is not None and args.n != 1: - parameters["n"] = args.n - if args.seed is not None: - parameters["seed"] = args.seed - if args.watermark is not None: - parameters["watermark"] = args.watermark - if args.prompt_extend is not None: - parameters["prompt_extend"] = args.prompt_extend - - try: - response = await AioMultiModalConversation.call( - api_key=api_key, - model=model_name, - messages=messages, - **parameters, - ) - except Exception as e: - raise RuntimeError( - f"Failed to call Wanx 2.6 image generation API: {str(e)}", - ) from e - - if response.status_code != 200 or not response.output: - raise RuntimeError(f"Wanx 2.6 image generation failed: {response}") - - results = [] - try: - if hasattr(response, "output") and response.output: - choices = getattr(response.output, "choices", []) - if choices: - message = getattr(choices[0], "message", {}) - content = getattr(message, "content", []) - if isinstance(content, list): - for item in content: - if isinstance(item, dict) and "image" in item: - results.append(item["image"]) - elif isinstance(content, str): - results.append(content) - elif isinstance(content, dict) and "image" in content: - results.append(content["image"]) - except Exception as e: - raise RuntimeError( - f"Failed to parse Wanx 2.6 API response: {str(e)}", - ) from e - - if not results: - raise RuntimeError(f"No image URLs found in response: {response}") - - if not request_id: - request_id = getattr(response, "request_id", None) or str( - uuid.uuid4(), - ) - - if trace_event: - trace_event.on_log( - "", - **{ - "step_suffix": "results", - "payload": { - "request_id": request_id, - "wanx26_image_generation_result": { - "status_code": response.status_code, - "results": results, - }, - }, - }, - ) - - return ImageGenerationWan26Output( - results=results, - request_id=request_id, - ) +# -*- coding: utf-8 -*- +import uuid +from typing import Any, Optional +from dashscope import AioMultiModalConversation +from mcp.server.fastmcp import Context +from pydantic import BaseModel, Field + +from agentscope_bricks.base.component import Component +from agentscope_bricks.utils.tracing_utils.wrapper import trace, TraceType +from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key +from agentscope_bricks.utils.tracing_utils import TracingUtil + + +class ImageGenerationWan26Input(BaseModel): + """ + Input schema for Wanx 2.6 text-to-image generation. + """ + + prompt: str = Field( + ..., + description="正向提示词，描述期望生成的图像内容，建议详细且清晰。超过800字符将被截断。", + ) + negative_prompt: Optional[str] = Field( + default=None, + description="反向提示词，描述不希望出现的内容，如低质量、模糊、文字等。超过500字符将被截断。", + ) + size: Optional[str] = Field( + default=None, + description="输出图像的分辨率。默认值是1280*1280，可不填。", + ) + prompt_extend: Optional[bool] = Field( + default=None, + description="是否开启 Prompt 智能改写。将使用大模型优化正向提示词。true: 开启（默认），false：不开启。", + ) + n: Optional[int] = Field( + default=1, + description="生成图片的数量。取值范围为1~4张 默认1", + ) + seed: Optional[int] = Field( + default=None, + description="随机种子，用于结果复现。", + ) + watermark: Optional[bool] = Field( + default=None, + description="是否添加阿里云水印，默认不添加。" + "false：默认值，不添加水印,true：添加水印。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context for " + "MCP internal use only, do not generate it.", + ) + + +class ImageGenerationWan26Output(BaseModel): + """ + Output schema for Wanx 2.6 text-to-image generation. + """ + + results: list[str] = Field( + title="Results", + description="生成的图片URL列表。", + ) + request_id: Optional[str] = Field( + default=None, + title="Request ID", + description="本次请求的唯一标识。", + ) + + +class ImageGenerationWan26( + Component[ImageGenerationWan26Input, ImageGenerationWan26Output], +): + """ + Wanx 2.6 Text-to-Image Generation Tool. + Uses the 'wan2.6-t2i' model from DashScope + to generate high-quality images from text. + """ + + name: str = "modelstudio_wanx26_image_generation" + description: str = ( + "[版本: wan2.6] 通义万相文生图模型（wanx2.6-t2i）。AI绘画服务，根据文本描述生成高质量图像，并返回图片URL。\n" + "新功能包括图像编辑和图文混合输出，满足更多样化的生成与集成需求。\n" + "支持自定义分辨率：图像面积介于 768×768 至 1440×1440 像素之间，" + "允许在该范围内自由调整宽高比（例如 768×2700）。\n" + ) + + @trace(trace_type=TraceType.AIGC, trace_name="wanx26_image_generation") + async def arun( + self, + args: ImageGenerationWan26Input, + **kwargs: Any, + ) -> ImageGenerationWan26Output: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") + + model_name = "wan2.6-t2i" + messages = [ + { + "role": "user", + "content": [{"text": args.prompt}], + }, + ] + + # Normalize watermark + if args.watermark is not None: + if isinstance(args.watermark, str): + args.watermark = args.watermark.strip().lower() in ( + "true", + "1", + ) + else: + args.watermark = bool(args.watermark) + + parameters = {} + if args.negative_prompt: + parameters["negative_prompt"] = args.negative_prompt + if args.size and args.size != "1024*1024": + parameters["size"] = args.size + if args.n is not None and args.n != 1: + parameters["n"] = args.n + if args.seed is not None: + parameters["seed"] = args.seed + if args.watermark is not None: + parameters["watermark"] = args.watermark + if args.prompt_extend is not None: + parameters["prompt_extend"] = args.prompt_extend + + try: + response = await AioMultiModalConversation.call( + api_key=api_key, + model=model_name, + messages=messages, + **parameters, + ) + except Exception as e: + raise RuntimeError( + f"Failed to call Wanx 2.6 image generation API: {str(e)}", + ) from e + + if response.status_code != 200 or not response.output: + raise RuntimeError(f"Wanx 2.6 image generation failed: {response}") + + results = [] + try: + if hasattr(response, "output") and response.output: + choices = getattr(response.output, "choices", []) + if choices: + message = getattr(choices[0], "message", {}) + content = getattr(message, "content", []) + if isinstance(content, list): + for item in content: + if isinstance(item, dict) and "image" in item: + results.append(item["image"]) + elif isinstance(content, str): + results.append(content) + elif isinstance(content, dict) and "image" in content: + results.append(content["image"]) + except Exception as e: + raise RuntimeError( + f"Failed to parse Wanx 2.6 API response: {str(e)}", + ) from e + + if not results: + raise RuntimeError(f"No image URLs found in response: {response}") + + if not request_id: + request_id = getattr(response, "request_id", None) or str( + uuid.uuid4(), + ) + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "results", + "payload": { + "request_id": request_id, + "wanx26_image_generation_result": { + "status_code": response.status_code, + "results": results, + }, + }, + }, + ) + + return ImageGenerationWan26Output( + results=results, + request_id=request_id, + ) diff --git a/src/agentscope_bricks/components/generations/image_generation_zimage.py b/src/agentscope_bricks/components/generations/image_generation_zimage.py new file mode 100644 index 0000000..76cf610 --- /dev/null +++ b/src/agentscope_bricks/components/generations/image_generation_zimage.py @@ -0,0 +1,162 @@ +# -*- coding: utf-8 -*- +import uuid +from typing import Any, Optional +from dashscope import AioMultiModalConversation +from mcp.server.fastmcp import Context +from pydantic import BaseModel, Field + +from agentscope_bricks.base.component import Component +from agentscope_bricks.utils.tracing_utils.wrapper import trace, TraceType +from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key +from agentscope_bricks.utils.tracing_utils import TracingUtil + + +class ZImageGenerationInput(BaseModel): + """ + Input schema for Z-Image text-to-image generation. + """ + + prompt: str = Field( + ..., + description="正向提示词，描述期望生成的图像内容，建议详细且清晰。超过800字符将被截断。", + ) + size: Optional[str] = Field( + default="1024*1536", + description="输出图像的分辨率。默认 1024*1536", + ) + prompt_extend: Optional[bool] = Field( + default=None, + description="是否开启 Prompt 智能改写。将使用大模型优化正向提示词。true: 开启，false：不开启（默认）。", + ) + seed: Optional[int] = Field( + default=None, + description="随机种子，用于结果复现。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context for MCP " + "internal use only, do not generate it.", + ) + + +class ZImageGenerationOutput(BaseModel): + """ + Output schema for Z-Image text-to-image generation. + """ + + results: list[str] = Field( + title="Results", + description="生成的图片URL列表。", + ) + request_id: Optional[str] = Field( + default=None, + title="Request ID", + description="本次请求的唯一标识。", + ) + + +class ZImageGeneration( + Component[ZImageGenerationInput, ZImageGenerationOutput], +): + """ + Z-Image Text-to-Image Generation Tool (based on z-image-turbo). + Uses the 'z-image-turbo' model from DashScope to + generate high-quality images from text prompts. + Supports custom resolution, negative prompts, batch generation, and more. + """ + + name: str = "modelstudio_z_image_generation" + description: str = ( + " 基于通义Z-Image大模型的智能图像生成服务，是一款轻量级文生图模型，" + "可快速生成图像，支持中英文字渲染，并灵活适配多种分辨率与宽高比例。" + ) + + @trace(trace_type=TraceType.AIGC, trace_name="z_image_generation") + async def arun( + self, + args: ZImageGenerationInput, + **kwargs: Any, + ) -> ZImageGenerationOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") + model_name = "z-image-turbo" + messages = [ + { + "role": "user", + "content": [{"text": args.prompt}], + }, + ] + parameters = {} + if args.size and args.size != "1024*1536": + parameters["size"] = args.size + if args.seed is not None: + parameters["seed"] = args.seed + if args.prompt_extend is not None: + parameters["prompt_extend"] = args.prompt_extend + + try: + response = await AioMultiModalConversation.call( + api_key=api_key, + model=model_name, + messages=messages, + **parameters, + ) + except Exception as e: + raise RuntimeError( + f"Failed to call Z-Image (z-image-turbo) API: {str(e)}", + ) from e + + if response.status_code != 200 or not response.output: + raise RuntimeError(f"Z-Image generation failed: {response}") + results = [] + try: + choices = getattr(response.output, "choices", []) + if choices: + message = getattr(choices[0], "message", {}) + content = getattr(message, "content", []) + if isinstance(content, list): + for item in content: + if isinstance(item, dict) and "image" in item: + results.append(item["image"]) + elif isinstance(content, str): + results.append(content) + elif isinstance(content, dict) and "image" in content: + results.append(content["image"]) + except Exception as e: + raise RuntimeError( + f"Failed to parse Z-Image API response: {str(e)}", + ) from e + + if not results: + raise RuntimeError( + f"No image URLs found in Z-Image response: {response}", + ) + + if not request_id: + request_id = getattr(response, "request_id", None) or str( + uuid.uuid4(), + ) + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "results", + "payload": { + "request_id": request_id, + "z_image_generation_result": { + "status_code": response.status_code, + "results": results, + }, + }, + }, + ) + + return ZImageGenerationOutput( + results=results, + request_id=request_id, + ) diff --git a/src/agentscope_bricks/components/generations/image_out_painting.py b/src/agentscope_bricks/components/generations/image_out_painting.py new file mode 100644 index 0000000..b4889bc --- /dev/null +++ b/src/agentscope_bricks/components/generations/image_out_painting.py @@ -0,0 +1,271 @@ +# -*- coding: utf-8 -*- +import os +import uuid +import asyncio +from http import HTTPStatus +from typing import Any, Optional, Dict + +import aiohttp +from mcp.server.fastmcp import Context +from pydantic import BaseModel, Field + +from agentscope_bricks.base.component import Component +from agentscope_bricks.utils.tracing_utils.wrapper import trace +from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key +from agentscope_bricks.utils.tracing_utils import TracingUtil + + +DASHSCOPE_API_BASE = "https://dashscope.aliyuncs.com/api/v1" + + +class ImageOutPaintingAutoInput(BaseModel): + """ + Input for auto-submit-and-fetch image out-painting task. + """ + + image_url: str = Field( + ..., + description="输入图像的公网可访问 URL。", + ) + angle: Optional[float] = Field( + default=None, + description="逆时针旋转角度，取值范围 [0, 359]。默认为 0（不旋转）。", + ) + output_ratio: Optional[str] = Field( + default=None, + description='目标宽高比，可选值：["", "1:1", "3:4", "4:3", "9:16", "16:9"]。' + '默认值为""，表示不设置输出图像的宽高比。', + ) + x_scale: Optional[float] = Field( + default=None, + description="水平方向扩展比例（居中扩展），默认 1.0。可以与 y_scale 搭配使用。取值范围 [1.0, 3.0]。" + "例如：输入图像分辨率为1000×1000（宽×高），x_scale=2.0，扩展后的图像分辨率为2000×1000（宽×高）。" + "保持高度不变，左右各添加500个像素。", + ) + y_scale: Optional[float] = Field( + default=None, + description="垂直方向扩展比例（居中扩展），默认 1.0。可以选择与 x_scale 搭配使用。取值范围 [1.0, 3.0]。" + "例如：输入图像分辨率为1000×1000（宽×高），y_scale=2.0，扩展后的图像分辨率为1000×2000（宽×高）。" + "保持宽度不变，上下各添加500个像素。", + ) + top_offset: Optional[float] = Field( + default=None, + description="在图像上方添加的像素数。默认值为0," + "需满足 top_offset + bottom_offset < 3 × 原图高度。" + "输入图像分辨率为1000×1000（宽×高），top_offset=500，扩展后的图像分辨率为1000×1500（宽×高）。" + "保持宽度不变，只在图像上方添加500个像素。", + ) + bottom_offset: Optional[float] = Field( + default=None, + description="在图像下方添加的像素数。默认值为0，" + "需满足 top_offset + bottom_offset < 3 × 原图高度。" + "例如：输入图像分辨率为1000×1000（宽×高），bottom_offset=500，扩展后的图像分辨率为1000×1500（宽×高）。" + "保持宽度不变，只在图像下方添加500个像素。", + ) + left_offset: Optional[float] = Field( + default=None, + description="在图像左侧添加的像素数。默认值为0，" + "需满足 left_offset + right_offset < 3 × 原图宽度。" + "例如：输入图像分辨率为1000×1000（宽×高），left_offset=500，扩展后的图像分辨率为1500×1000（宽×高）。" + "保持高度不变，只在图像左侧添加500个像素。", + ) + right_offset: Optional[float] = Field( + default=None, + description="在图像右侧添加的像素数。默认值为0，" + "需满足 left_offset + right_offset < 3 × 原图宽度。" + "例如：输入图像分辨率为1000×1000（宽×高），right_offset=500，扩展后的图像分辨率为1500×1000（宽×高）。" + "保持高度不变，只在图像右侧添加500个像素。", + ) + best_quality: Optional[bool] = Field( + default=None, + description="是否启用最佳质量模式。默认 false（速度优先），设为 true 可提升细节但耗时增加。", + ) + limit_image_size: Optional[bool] = Field( + default=None, + description="是否限制输出图像大小（≤5MB）。默认 true，建议保持开启。" + "模型生成的图像需要经过一层安全过滤后才能输出，当前不支持大于10M的图像处理。", + ) + add_watermark: Optional[bool] = Field( + default=None, + description="是否添加水印,True：默认值,添加水印,False：不添加水印。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context containing " + "headers for mcp only, don't generate it", + ) + + +class ImageOutPaintingAutoOutput(BaseModel): + output_image_url: str = Field( + ..., + description="扩图后生成的图像公网 URL（PNG/JPG 等格式），有效期 24 小时。", + ) + task_id: str = Field( + ..., + description="异步任务的唯一标识符。", + ) + request_id: Optional[str] = Field( + default=None, + description="请求 ID，用于日志追踪。", + ) + + +class ImageOutPaintingAuto( + Component[ImageOutPaintingAutoInput, ImageOutPaintingAutoOutput], +): + name: str = "modelstudio_image_out_painting_auto" + description: str = ( + "图像画面扩展（扩图）同步自动执行工具。\n" + "提交扩图任务并内部轮询结果，直接返回扩图后的图像 URL。\n" + "无需手动查询任务状态，适合需要端到端结果的场景。" + ) + + @trace(trace_type="AIGC", trace_name="image_out_painting_auto") + async def arun( + self, + args: ImageOutPaintingAutoInput, + **kwargs: Any, + ) -> ImageOutPaintingAutoOutput: + trace_event = kwargs.pop("trace_event", None) + request_id = TracingUtil.get_request_id() + + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") + + # 构建 parameters（仅非 None 值） + parameters: Dict[str, Any] = {} + for field in [ + "angle", + "output_ratio", + "x_scale", + "y_scale", + "top_offset", + "bottom_offset", + "left_offset", + "right_offset", + "best_quality", + "limit_image_size", + "add_watermark", + ]: + value = getattr(args, field) + if value is not None: + parameters[field] = value + + headers = { + "Authorization": f"Bearer {api_key}", + "X-DashScope-Async": "enable", + "Content-Type": "application/json", + } + + payload = { + "model": "image-out-painting", + "input": {"image_url": args.image_url}, + "parameters": parameters, + } + + # Step 1: Submit task + async with aiohttp.ClientSession() as session: + async with session.post( + f"{DASHSCOPE_API_BASE}/services/aigc/image2image/out-painting", + headers=headers, + json=payload, + ) as resp: + status_code = resp.status + response_json = await resp.json() + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "submit", + "payload": { + "request_id": request_id, + "response": response_json, + "status_code": status_code, + }, + }, + ) + + if status_code != HTTPStatus.OK or "output" not in response_json: + error_msg = response_json.get("message", "Unknown error") + raise RuntimeError( + f"Failed to submit out-painting task: {error_msg} (code: {status_code})", # noqa + ) + + task_id = response_json["output"]["task_id"] + request_id = ( + response_json.get("request_id") or request_id or str(uuid.uuid4()) + ) + + # Step 2: Poll until completion + max_retries = 60 # 最多等待 2 分钟（60 * 2s） + retry_interval = 2 # 每 2 秒查询一次 + + fetch_headers = {"Authorization": f"Bearer {api_key}"} + + for attempt in range(max_retries): + await asyncio.sleep(retry_interval) + + async with aiohttp.ClientSession() as session: + async with session.get( + f"{DASHSCOPE_API_BASE}/tasks/{task_id}", + headers=fetch_headers, + ) as resp: + fetch_status = resp.status + fetch_response = await resp.json() + + if fetch_status != HTTPStatus.OK or "output" not in fetch_response: + error_msg = fetch_response.get( + "message", + "Unknown fetch error", + ) + raise RuntimeError( + f"Failed to poll task: {error_msg} (code: {fetch_status})", + ) + + output = fetch_response["output"] + task_status = output["task_status"] + + if task_status == "SUCCEEDED": + output_image_url = output["output_image_url"] + final_request_id = ( + fetch_response.get("request_id") or request_id + ) + + if trace_event: + trace_event.on_log( + "", + **{ + "step_suffix": "success", + "payload": { + "output_image_url": output_image_url, + "request_id": final_request_id, + }, + }, + ) + + return ImageOutPaintingAutoOutput( + output_image_url=output_image_url, + task_id=task_id, + request_id=final_request_id, + ) + + elif task_status in ("FAILED", "CANCELED"): + error_msg = output.get( + "message", + "Task failed without details", + ) + raise RuntimeError( + f"Out-painting task failed: {error_msg} (task_id: {task_id})", # noqa + ) + + # else: PENDING / RUNNING → continue polling + + # Timeout + raise TimeoutError( + f"Out-painting task did not complete within {max_retries * retry_interval} seconds " # noqa + f"(task_id: {task_id}). Current status may still be PENDING/RUNNING.", # noqa + ) diff --git a/src/agentscope_bricks/components/generations/image_text_interleave_generation_wan26.py b/src/agentscope_bricks/components/generations/image_text_interleave_generation_wan26.py new file mode 100644 index 0000000..b4e0e42 --- /dev/null +++ b/src/agentscope_bricks/components/generations/image_text_interleave_generation_wan26.py @@ -0,0 +1,198 @@ +# -*- coding: utf-8 -*- +import uuid +import json +from http import HTTPStatus +from typing import Any, Optional, Dict, AsyncGenerator +import aiohttp +from mcp.server.fastmcp import Context +from pydantic import BaseModel, Field + +from agentscope_bricks.base.component import Component +from agentscope_bricks.utils.tracing_utils.wrapper import trace +from agentscope_bricks.utils.api_key_util import ApiNames, get_api_key +from agentscope_bricks.utils.tracing_utils import TracingUtil + + +DASHSCOPE_API_BASE = "https://dashscope.aliyuncs.com/api/v1" + + +class WanImageInterleaveGenerationInput(BaseModel): + """ + Input model for Alibaba Cloud + Wan 2.6 Image Interleaved (Text + Image) Generation. + """ + + prompt: str = Field( + ..., + description="用户输入的文本指令，例如 '给我一个3张图辣椒炒肉教程'。", + ) + negative_prompt: Optional[str] = Field( + default=None, + description="反向提示词，描述不希望出现的内容，如低质量、模糊、文字等。", + ) + image: Optional[str] = Field( + default=None, + description="可选的参考图像 URL，图片和prompt要有关系，否则会被忽略。", + ) + max_images: Optional[int] = Field( + default=5, + description="期望生成的最大图像数量取值范围：1～5，默认值为 5,该参数仅代表“数量上限”。" + "实际生成的图像数量由模型推理决定，可能会少于设定值。", + ) + size: Optional[str] = Field( + default="1280*1280", + description="输出图像的分辨率。默认值是1280*1280，可不填。", + ) + watermark: Optional[bool] = Field( + default=None, + description="是否添加水印,false：默认值，不添加水印,true：添加水印。", + ) + seed: Optional[int] = Field( + default=None, + description="随机种子，用于结果可复现。", + ) + ctx: Optional[Context] = Field( + default=None, + description="HTTP request context containing " + "headers for mcp only, don't generate it", + ) + + +class WanImageInterleaveGenerationOutput(BaseModel): + full_text: str = Field( + ..., + description="模型生成的完整文本内容（不含图片占位符）。", + ) + image_urls: list[str] = Field( + ..., + description="按顺序生成的图像公网 URL 列表。", + ) + request_id: Optional[str] = Field( + default=None, + description="请求唯一 ID，用于日志追踪。", + ) + + +class WanImageInterleaveGeneration( + Component[ + WanImageInterleaveGenerationInput, + WanImageInterleaveGenerationOutput, + ], +): + name: str = "modelstudio_wan_text_image_interleave_generation" + description: str = ( + "[版本: wan2.6] 通义万相图文混排生成工具（wan2.6-image），支持文本+图像混合生成。\n" + "支持传入最多1张参考图用于风格/背景引导。" + ) + + @trace( + trace_type="AIGC", + trace_name="wan_image_interleave_generation_stream", + ) + async def astream( + self, + args: WanImageInterleaveGenerationInput, + **kwargs: Any, + ) -> AsyncGenerator[Dict[str, Any], None]: + try: + api_key = get_api_key(ApiNames.dashscope_api_key, **kwargs) + except AssertionError: + raise ValueError("Please set valid DASHSCOPE_API_KEY!") + + content: list[Dict[str, str]] = [{"text": args.prompt}] + if args.image: + content.append({"image": args.image}) + parameters = { + "enable_interleave": True, # 必须为 true + "stream": True, # 启用流式 + "max_images": args.max_images, + "size": args.size, + "watermark": args.watermark, + } + + # 可选参数：仅当非 None 时传入 + if args.negative_prompt is not None: + parameters["negative_prompt"] = args.negative_prompt + if args.seed is not None: + parameters["seed"] = args.seed + + payload = { + "model": "wan2.6-image", + "input": { + "messages": [ + { + "role": "user", + "content": content, + }, + ], + }, + "parameters": parameters, + } + + headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + "X-DashScope-Sse": "enable", + } + + async with aiohttp.ClientSession() as session: + async with session.post( + f"{DASHSCOPE_API_BASE}/services/aigc/multimodal-generation/generation", # noqa + headers=headers, + json=payload, + ) as resp: + if resp.status != HTTPStatus.OK: + error_text = await resp.text() + raise RuntimeError(f"SSE request failed: {error_text}") + + async for line_bytes in resp.content: + line = line_bytes.decode("utf-8").strip() + if not line or not line.startswith("data:"): + continue + + data_str = line[5:].strip() + if data_str == "[DONE]": + break + + try: + chunk = json.loads(data_str) + contents = chunk["output"]["choices"][0]["message"][ + "content" + ] + for item in contents: + if item.get("type") == "text": + yield {"type": "text", "value": item["text"]} + elif item.get("type") == "image": + img_url = item.get("image") + if isinstance(img_url, str): + yield {"type": "image", "value": img_url} + except ( + KeyError, + IndexError, + TypeError, + json.JSONDecodeError, + ): + continue + + @trace(trace_type="AIGC", trace_name="wan_image_interleave_generation") + async def arun( + self, + args: WanImageInterleaveGenerationInput, + **kwargs: Any, + ) -> WanImageInterleaveGenerationOutput: + full_text = "" + image_urls: list[str] = [] + request_id = TracingUtil.get_request_id() or str(uuid.uuid4()) + + # 复用 astream 逻辑来聚合结果（避免重复代码） + async for chunk in self.astream(args, **kwargs): + if chunk["type"] == "text": + full_text += chunk["value"] + elif chunk["type"] == "image": + image_urls.append(chunk["value"]) + + return WanImageInterleaveGenerationOutput( + full_text=full_text, + image_urls=image_urls, + request_id=request_id, + )