agentscope-runtime 0.2.0b2__py3-none-any.whl → 1.0.0b1__py3-none-any.whl

agentscope_runtime/tools/generations/image_generation_wan25.py

@@ -0,0 +1,201 @@
+# -*- coding: utf-8 -*-
+# pylint:disable=no-else-break, too-many-branches, abstract-method
+
+import asyncio
+import os
+import time
+import uuid
+from http import HTTPStatus
+from typing import Any, Optional
+
+from dashscope.aigc.image_synthesis import AioImageSynthesis
+from mcp.server.fastmcp import Context
+from pydantic import BaseModel, Field
+
+from ..base import Tool
+from ..utils.api_key_util import ApiNames, get_api_key
+from ...engine.tracing import trace, TracingUtil
+
+
+class ImageGenInput(BaseModel):
+    """
+    Text-to-Image Input.
+    """
+
+    prompt: str = Field(
+        ...,
+        description="正向提示词，用来描述生成图像中期望包含的元素和视觉特点,超过800自动截断",
+    )
+    size: Optional[str] = Field(
+        default=None,
+        description="输出图像的分辨率。默认值是1280*1280，可不填。",
+    )
+    negative_prompt: Optional[str] = Field(
+        default=None,
+        description="反向提示词，用来描述不希望在画面中看到的内容，可以对画面进行限制，超过500个字符自动截断",
+    )
+    prompt_extend: Optional[bool] = Field(
+        default=None,
+        description="是否开启prompt智能改写，开启后使用大模型对输入prompt进行智能改写",
+    )
+    n: Optional[int] = Field(
+        default=1,
+        description="生成图片的数量。取值范围为1~4张 默认1",
+    )
+    watermark: Optional[bool] = Field(
+        default=None,
+        description="是否添加水印，默认不设置",
+    )
+    ctx: Optional[Context] = Field(
+        default=None,
+        description="HTTP request context containing headers for mcp only, "
+        "don't generate it",
+    )
+
+
+class ImageGenOutput(BaseModel):
+    """
+    Text-to-Image Output.
+    """
+
+    results: list[str] = Field(title="Results", description="输出图片url 列表")
+    request_id: Optional[str] = Field(
+        default=None,
+        title="Request ID",
+        description="请求ID",
+    )
+
+
+class ImageGenerationWan25(Tool[ImageGenInput, ImageGenOutput]):
+    """
+    Text-to-Image Call.
+    """
+
+    name: str = "modelstudio_image_gen_wan25"
+    description: str = "AI绘画（图像生成）服务，输入文本描述和图像分辨率，返回根据文本信息绘制的图片URL。"
+
+    @trace(trace_type="AIGC", trace_name="image_generation_wan25")
+    async def arun(self, args: ImageGenInput, **kwargs: Any) -> ImageGenOutput:
+        """Modelstudio Images generation from text prompts
+
+        This method wrap DashScope's ImageSynthesis service to generate images
+        based on text descriptions. It supports various image sizes and can
+        generate multiple images in a single request.
+
+        Args:
+            args: ImageGenInput containing the prompt, size, and number of
+                images to generate.
+            **kwargs: Additional keyword arguments including:
+                - request_id: Optional request ID for tracking
+                - trace_event: Optional trace event for logging
+                - model_name: Model name to use (defaults to wan2.2-t2i-flash)
+                - api_key: DashScope API key for authentication
+
+        Returns:
+            ImageGenOutput containing the list of generated image URLs and
+            request ID.
+
+        Raises:
+            ValueError: If DASHSCOPE_API_KEY is not set or invalid.
+        """
+
+        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
+
+        model_name = kwargs.get(
+            "model_name",
+            os.getenv("IMAGE_GENERATION_MODEL_NAME", "wan2.5-t2i-preview"),
+        )
+
+        parameters = {}
+        if args.size:
+            parameters["size"] = args.size
+        if args.prompt_extend is not None:
+            parameters["prompt_extend"] = args.prompt_extend
+        if args.n is not None:
+            parameters["n"] = args.n
+        if args.watermark is not None:
+            parameters["watermark"] = args.watermark
+
+        task_response = await AioImageSynthesis.async_call(
+            model=model_name,
+            api_key=api_key,
+            prompt=args.prompt,
+            negative_prompt=args.negative_prompt,
+            **parameters,
+        )
+
+        if (
+            task_response.status_code != HTTPStatus.OK
+            or not task_response.output
+        ):
+            raise RuntimeError(f"Failed to submit task: {task_response}")
+
+        # 2. Loop to asynchronously query task status
+        max_wait_time = 300  # 5 minutes timeout
+        poll_interval = 2  # 2 seconds polling interval
+        start_time = time.time()
+
+        while True:
+            # Asynchronous wait
+            await asyncio.sleep(poll_interval)
+
+            # Query task result
+            res = await AioImageSynthesis.fetch(
+                api_key=api_key,
+                task=task_response,
+            )
+
+            if (
+                res.status_code != HTTPStatus.OK
+                or not res.output
+                or (
+                    hasattr(res.output, "task_status")
+                    and res.output.task_status in ["FAILED", "CANCELED"]
+                )
+            ):
+                raise RuntimeError(f"Failed to fetch result: {res}")
+
+            # Check if task is completed
+            if res.status_code == HTTPStatus.OK:
+                if hasattr(res.output, "task_status"):
+                    if res.output.task_status == "SUCCEEDED":
+                        break
+                    elif res.output.task_status in ["FAILED", "CANCELED"]:
+                        raise RuntimeError(f"Failed to generate: {res}")
+                else:
+                    # If no task_status field, consider task completed
+                    break
+
+            # Timeout check
+            if time.time() - start_time > max_wait_time:
+                raise TimeoutError(
+                    f"Image generation timeout after {max_wait_time}s",
+                )
+
+        if request_id == "":
+            request_id = (
+                res.request_id if res.request_id else str(uuid.uuid4())
+            )
+
+        if trace_event:
+            trace_event.on_log(
+                "",
+                **{
+                    "step_suffix": "results",
+                    "payload": {
+                        "request_id": request_id,
+                        "image_query_result": res,
+                    },
+                },
+            )
+        results = []
+        if res.status_code == HTTPStatus.OK:
+            for result in res.output.results:
+                results.append(result.url)
+        return ImageGenOutput(results=results, request_id=request_id)

agentscope_runtime/tools/generations/image_style_repaint.py

@@ -0,0 +1,208 @@
+# -*- coding: utf-8 -*-
+# pylint:disable=abstract-method, deprecated-module, wrong-import-order
+# pylint:disable=redefined-builtin
+
+import asyncio
+import os
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from http import HTTPStatus
+from typing import Any, Optional
+
+from dashscope.client.base_api import BaseAsyncApi
+from dashscope.utils.oss_utils import check_and_upload_local
+from mcp.server.fastmcp import Context
+from pydantic import BaseModel, Field
+
+from ..base import Tool
+from ..utils.api_key_util import get_api_key, ApiNames
+from ...engine.tracing import trace, TracingUtil
+
+
+class ImageStyleRepaintInput(BaseModel):
+    """
+    Portrait Style Repaint Input
+    """
+
+    image_url: str = Field(
+        ...,
+        description="输入图像的URL地址。",
+    )
+
+    style_index: int = Field(
+        ...,
+        description="人像风格类型索引值，当前支持以下风格：-1：参考上传图像风格, "
+        "0：复古漫画, 1：3D童话, 2：二次元, 3：小清新, 4：未来科技, "
+        "5：国画古风, 6：将军百战, 7：炫彩卡通, 8：清雅国风, 9：喜迎新年。",
+    )
+
+    style_ref_url: Optional[str] = Field(
+        default=None,
+        description="风格参考图像的URL地址。当参数style_index等于-1时，必须传入，" "其他风格无需传入。",
+    )
+
+    watermark: Optional[bool] = Field(
+        default=None,
+        description="是否添加水印，默认不设置",
+    )
+
+    ctx: Optional[Context] = Field(
+        default=None,
+        description="HTTP request context containing headers for mcp only, "
+        "don't generate it",
+    )
+
+
+class ImageStyleRepaintOutput(BaseModel):
+    """
+    Portrait Style Repaint Output
+    """
+
+    results: list[str] = Field(title="Results", description="输出图片url 列表")
+    request_id: Optional[str] = Field(
+        default=None,
+        title="Request ID",
+        description="请求ID",
+    )
+
+
+class ImageStyleRepaint(
+    Tool[ImageStyleRepaintInput, ImageStyleRepaintOutput],
+):
+    """
+    Portrait Style Repaint
+    """
+
+    name: str = "modelstudio_image_style_repaint"
+    description: str = "人像风格重绘服务，输入原始图像和风格数据(索引或参考图像），返回重绘后的图像。"
+
+    def __init__(self, name: str = None, description: str = None):
+        super().__init__(name=name, description=description)
+        # Create thread pool to execute synchronous BaseAsyncApi calls
+        self._executor = ThreadPoolExecutor(
+            max_workers=10,
+            thread_name_prefix="StyleRepaint",
+        )
+
+    @trace(trace_type="AIGC", trace_name="image_style_repaint")
+    async def arun(
+        self,
+        args: ImageStyleRepaintInput,
+        **kwargs: Any,
+    ) -> ImageStyleRepaintOutput:
+        """Modelstudio Image Style Repaint
+
+        This method wrap DashScope's ImageStyleRepaint service to generate
+        images based on image url and style index (or style reference image
+        url).
+
+        Args:
+            args: ImageStyleRepaintInput containing the image_url,
+                style_index, and style_ref_url.
+            **kwargs: Additional keyword arguments including:
+                - request_id: Optional request ID for tracking
+                - trace_event: Optional trace event for logging
+                - model_name: Model name to use (defaults to wanx2.1-t2i-turbo)
+                - api_key: DashScope API key for authentication
+
+        Returns:
+            ImageStyleRepaintOutput containing the list of generated image
+            URLs and request ID.
+
+        Raises:
+            ValueError: If DASHSCOPE_API_KEY is not set or invalid.
+        """
+
+        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
+
+        model_name = kwargs.get(
+            "model_name",
+            os.getenv(
+                "IMAGE_STYLE_REPAINT_MODEL_NAME",
+                "wanx-style-repaint-v1",
+            ),
+        )
+
+        has_uploaded = False
+
+        image_url = args.image_url
+        if args.image_url:
+            uploaded, image_url, _ = check_and_upload_local(
+                model=model_name,
+                content=args.image_url,
+                api_key=api_key,
+            )
+            has_uploaded = True if uploaded is True else has_uploaded
+
+        style_ref_url = args.style_ref_url
+        if args.style_ref_url:
+            uploaded, style_ref_url = check_and_upload_local(
+                model=model_name,
+                content=args.style_ref_url,
+                api_key=api_key,
+            )
+            has_uploaded = True if uploaded is True else has_uploaded
+
+        kwargs = {}
+        if has_uploaded is True:
+            headers = {"X-DashScope-OssResourceResolve": "enable"}
+            kwargs["headers"] = headers
+
+        # 🔄 Put BaseAsyncApi.call into thread pool to avoid blocking
+        # event loop
+        def _sync_style_repaint_call() -> Any:
+            input = {
+                "image_url": image_url,
+                "style_index": args.style_index,
+                "style_ref_url": style_ref_url,
+            }
+            if args.watermark is not None:
+                input["watermark"] = args.watermark
+            return BaseAsyncApi.call(
+                model=model_name,
+                input=input,
+                task_group="aigc",
+                task="image-generation",
+                function="generation",
+                **kwargs,
+            )
+
+        # Execute synchronous calls asynchronously in thread pool
+        res = await asyncio.get_event_loop().run_in_executor(
+            self._executor,
+            _sync_style_repaint_call,
+        )
+
+        if res.status_code != HTTPStatus.OK or not res.output:
+            raise RuntimeError(f"Failed to generate image: {res}")
+
+        if request_id == "":
+            request_id = (
+                res.request_id if res.request_id else str(uuid.uuid4())
+            )
+
+        if trace_event:
+            trace_event.on_log(
+                "",
+                **{
+                    "step_suffix": "results",
+                    "payload": {
+                        "request_id": request_id,
+                        "image_query_result": res,
+                    },
+                },
+            )
+
+        results = []
+        if res.status_code == HTTPStatus.OK:
+            for result in res.output.get("results"):
+                if result.get("url"):
+                    results.append(result.get("url"))
+
+        return ImageStyleRepaintOutput(results=results, request_id=request_id)

agentscope_runtime/tools/generations/image_to_video.py

@@ -0,0 +1,233 @@
+# -*- coding: utf-8 -*-
+# pylint:disable=abstract-method, deprecated-module, wrong-import-order
+# pylint:disable=no-else-break, too-many-branches
+
+import asyncio
+import os
+import time
+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 ..base import Tool
+from ..utils.api_key_util import get_api_key, ApiNames
+from ...engine.tracing import trace, TracingUtil
+
+
+class ImageToVideoInput(BaseModel):
+    """
+    Image to video generation input model
+    """
+
+    image_url: str = Field(
+        ...,
+        description="输入图像，支持公网URL、Base64编码或本地文件路径",
+    )
+    prompt: Optional[str] = Field(
+        default=None,
+        description="正向提示词，用来描述生成视频中期望包含的元素和视觉特点",
+    )
+    negative_prompt: Optional[str] = Field(
+        default=None,
+        description="反向提示词，用来描述不希望在视频画面中看到的内容",
+    )
+    template: Optional[str] = Field(
+        default=None,
+        description="视频特效模板，可选值：squish（解压捏捏）、flying（魔法悬浮）、carousel（时光木马）等",
+    )
+    resolution: Optional[str] = Field(
+        default=None,
+        description="视频分辨率，默认不设置",
+    )
+    duration: Optional[int] = Field(
+        default=None,
+        description="视频生成时长，单位为秒，通常为5秒",
+    )
+    prompt_extend: Optional[bool] = Field(
+        default=None,
+        description="是否开启prompt智能改写，开启后使用大模型对输入prompt进行智能改写",
+    )
+    watermark: Optional[bool] = Field(
+        default=None,
+        description="是否添加水印，默认不设置",
+    )
+    ctx: Optional[Context] = Field(
+        default=None,
+        description="HTTP request context containing headers for mcp only, "
+        "don't generate it",
+    )
+
+
+class ImageToVideoOutput(BaseModel):
+    """
+    Image to video generation output model
+    """
+
+    video_url: str = Field(
+        title="Video URL",
+        description="输出的视频url",
+    )
+    request_id: Optional[str] = Field(
+        default=None,
+        title="Request ID",
+        description="请求ID",
+    )
+
+
+class ImageToVideo(Tool[ImageToVideoInput, ImageToVideoOutput]):
+    """
+    Image to video generation service that converts images into videos
+    using DashScope's VideoSynthesis API.
+    """
+
+    name: str = "modelstudio_image_to_video"
+    description: str = (
+        "通义万相-图生视频模型根据首帧图像和文本提示词，生成时长为5秒的无声视频。"
+        "同时支持特效模板，可添加“魔法悬浮”、“气球膨胀”等效果，适用于创意视频制作、娱乐特效展示等场景。"
+    )
+
+    @trace(trace_type="AIGC", trace_name="image_to_video")
+    async def arun(
+        self,
+        args: ImageToVideoInput,
+        **kwargs: Any,
+    ) -> ImageToVideoOutput:
+        """
+        Generate video from image using DashScope VideoSynthesis
+
+        This method wraps DashScope's VideoSynthesis service to generate videos
+        based on input images. It uses async call pattern for better
+        performance and supports polling for task completion.
+
+        Args:
+            args: ImageToVideoInput containing required image_url and optional
+                  parameters
+            **kwargs: Additional keyword arguments including:
+                - request_id: Optional request ID for tracking
+                - model_name: Model name to use (defaults to wan2.2-i2v-flash)
+                - api_key: DashScope API key for authentication
+
+        Returns:
+            ImageToVideoOutput containing the generated video URL
+                and request ID
+
+        Raises:
+            ValueError: If DASHSCOPE_API_KEY is not set or invalid
+            TimeoutError: If video generation takes too long
+            RuntimeError: If video generation fails
+        """
+        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
+
+        model_name = kwargs.get(
+            "model_name",
+            os.getenv("IMAGE_TO_VIDEO_MODEL_NAME", "wan2.2-i2v-flash"),
+        )
+
+        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.watermark is not None:
+            parameters["watermark"] = args.watermark
+
+        # Create AioVideoSynthesis instance
+        aio_video_synthesis = AioVideoSynthesis()
+
+        # Submit async task
+        task_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,
+            template=args.template,
+            **parameters,
+        )
+
+        if (
+            task_response.status_code != HTTPStatus.OK
+            or not task_response.output
+            or task_response.output.task_status in ["FAILED", "CANCELED"]
+        ):
+            raise RuntimeError(f"Failed to submit task: {task_response}")
+
+        # Poll for task completion using async methods
+        max_wait_time = 600  # 10 minutes timeout for video generation
+        poll_interval = 5  # 5 seconds polling interval
+        start_time = time.time()
+
+        while True:
+            # Wait before polling
+            await asyncio.sleep(poll_interval)
+
+            # Fetch task result using async method
+            res = await aio_video_synthesis.fetch(
+                api_key=api_key,
+                task=task_response,
+            )
+
+            if (
+                res.status_code != HTTPStatus.OK
+                or not res.output
+                or res.output.task_status in ["FAILED", "CANCELED"]
+            ):
+                raise RuntimeError(f"Failed to fetch result: {res}")
+
+            # Check task completion status
+            if res.status_code == HTTPStatus.OK:
+                if hasattr(res.output, "task_status"):
+                    if res.output.task_status == "SUCCEEDED":
+                        break
+                    elif res.output.task_status in ["FAILED", "CANCELED"]:
+                        raise RuntimeError(f"Failed to generate: {res}")
+                else:
+                    # If no task_status field, assume completed
+                    break
+
+            # Check timeout
+            if time.time() - start_time > max_wait_time:
+                raise TimeoutError(
+                    f"Video generation timeout after {max_wait_time}s",
+                )
+
+        # Handle request ID
+        if not request_id:
+            request_id = (
+                res.request_id if res.request_id else str(uuid.uuid4())
+            )
+
+        # Log trace event if provided
+        if trace_event:
+            trace_event.on_log(
+                "",
+                **{
+                    "step_suffix": "results",
+                    "payload": {
+                        "request_id": request_id,
+                        "image_to_video_result": res,
+                    },
+                },
+            )
+
+        # Extract video URL from response
+        if res.status_code == HTTPStatus.OK:
+            video_url = res.output.video_url
+            return ImageToVideoOutput(
+                video_url=video_url,
+                request_id=request_id,
+            )
+        else:
+            raise RuntimeError(f"Failed to get video URL: {res.message}")