diff --git a/src/agentscope_bricks/components/__init__.py b/src/agentscope_bricks/components/__init__.py index 333ba7e..23fe223 100644 --- a/src/agentscope_bricks/components/__init__.py +++ b/src/agentscope_bricks/components/__init__.py @@ -15,9 +15,6 @@ 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, ) @@ -75,7 +72,26 @@ from agentscope_bricks.components.generations.fetch_wan import WanVideoFetch from agentscope_bricks.components.generations.qwen_image_edit_new import ( QwenImageEditNew, -) # noqa +) +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): @@ -92,7 +108,14 @@ class McpServerMeta(BaseModel): mcp_server_metas: Dict[str, McpServerMeta] = { "modelstudio_wan_image": McpServerMeta( instructions="基于通义万相大模型的智能图像生成服务，提供高质量的图像处理和编辑功能", - components=[ImageGeneration, ImageEdit, ImageStyleRepaint], + components=[ + ImageGeneration, + ImageEdit, + ImageStyleRepaint, + ImageOutPaintingSubmit, + ImageOutPaintingFetch, + ImageOutPaintingAuto, + ], ), "modelstudio_wan_video": McpServerMeta( instructions="基于通义万相大模型提供AI视频生成服务，支持文本到视频、图像到视频和语音到视频的多模态生成功能", @@ -103,6 +126,8 @@ class McpServerMeta(BaseModel): ImageToVideoFetch, SpeechToVideoSubmit, SpeechToVideoFetch, + ImageToVideoByFirstAndLastFrameWan22Submit, + WanVideoFetch, ], ), "modelstudio_wan25_media": McpServerMeta( @@ -130,7 +155,7 @@ class McpServerMeta(BaseModel): ), "modelstudio_speech_to_text": McpServerMeta( instructions="录音文件的语音识别服务，支持多种音频格式的语音转文字功能", - components=[SpeechToText, MultichannelSpeechToText], + components=[SpeechToText], ), "modelstudio_qwen_text_to_speech": McpServerMeta( instructions="基于通义千问大模型的语音合成服务，支持多种语言语音合成功能", @@ -143,6 +168,15 @@ class McpServerMeta(BaseModel): 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..a60b371 --- /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..12ddda3 --- /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..d1c2bdf 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 @@ -21,7 +21,7 @@ class ImageToVideoWan26SubmitInput(BaseModel): image_url: str = Field( ..., - description="输入图像，支持公网URL、Base64编码或本地文件路径", + description="输入图像，支持公网URL、Base64编码", ) prompt: Optional[str] = Field( default=None, @@ -33,15 +33,12 @@ class ImageToVideoWan26SubmitInput(BaseModel): ) audio_url: Optional[str] = Field( default=None, - description="自定义音频文件的公网URL。参数优先级：audio_url > audio。", - ) - audio: Optional[bool] = Field( - default=None, - description="是否自动生成配音。仅在 audio_url 未提供时生效。", + description="自定义音频文件的公网URL。参数优先级：audio_url > audio。" + "若不提供audio_url ，模型将根据视频内容自动生成匹配的背景音乐或音效。", ) template: Optional[str] = Field( default=None, - description="视频特效模板，如：squish（解压捏捏）、flying（魔法悬浮）、carousel（时光木马）等。", + description="视频特效模板，如：flying，表示使用“魔法悬浮”特效等。", ) resolution: Optional[str] = Field( default=None, @@ -64,7 +61,7 @@ class ImageToVideoWan26SubmitInput(BaseModel): ) watermark: Optional[bool] = Field( default=None, - description="是否在视频中添加水印（如“AI生成”标识）。默认不添加。", + description="是否添加水印,false：默认值，不添加水印,true：添加水印。", ) seed: Optional[int] = Field( default=None, @@ -134,8 +131,6 @@ async def arun( # 构建 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: 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..8b8bca0 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 @@ -30,12 +30,7 @@ class TextToVideoWan26SubmitInput(BaseModel): 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 生效。", + "若不提供audio_url ，模型将根据视频内容自动生成匹配的背景音乐或音效。", ) size: Optional[str] = Field( default=None, @@ -57,7 +52,7 @@ class TextToVideoWan26SubmitInput(BaseModel): ) watermark: Optional[bool] = Field( default=None, - description="是否添加水印，默认不设置", + description="是否在视频中添加水印,false：默认值，不添加水印,true：添加水印。", ) seed: Optional[int] = Field( default=None, @@ -128,8 +123,6 @@ async def arun( ) parameters = {} - if args.audio is not None: - parameters["audio"] = args.audio if args.size: parameters["size"] = args.size if args.duration is not None: 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..97224b8 --- /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..e68f767 100644 --- a/src/agentscope_bricks/components/generations/image_generation_wan26.py +++ b/src/agentscope_bricks/components/generations/image_generation_wan26.py @@ -42,7 +42,7 @@ class ImageGenerationWan26Input(BaseModel): ) watermark: Optional[bool] = Field( default=None, - description="是否添加阿里云水印，默认不添加。", + description="是否添加水印,false：默认值，不添加水印,true：添加水印。", ) ctx: Optional[Context] = Field( default=None, 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..1cbb2ef --- /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..bc5578b --- /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..e2e9d0b --- /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, + )