Files

763 lines
23 KiB
Python

# -*- coding: utf-8 -*-
# pylint: disable=too-many-return-statements,too-many-branches
# pylint: disable=too-many-statements,too-many-locals
"""Qwen-Image image generation and editing tools."""
import asyncio
import base64
import logging
import mimetypes
import threading
import time
from pathlib import Path
from typing import List
import httpx
from agentscope.message import DataBlock, TextBlock, URLSource
from agentscope.message import ToolResultState
from agentscope.tool import ToolChunk
from pineagents.constant import DEFAULT_MEDIA_DIR
from pineagents.plugins import get_tool_config
logger = logging.getLogger(__name__)
# Thread lock to protect dashscope global base_http_api_url setting
_DASHSCOPE_LOCK = threading.Lock()
_DEFAULT_ENDPOINT = "https://dashscope.aliyuncs.com/api/v1"
_DEFAULT_TIMEOUT = 120.0
_IMAGE_MIME_TYPES = {
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".webp": "image/webp",
}
_VALID_MODELS_GENERATE = {
"qwen-image-2.0-pro",
"qwen-image-2.0-pro-2026-04-22",
"qwen-image-2.0-pro-2026-03-03",
"qwen-image-2.0",
"qwen-image-2.0-2026-03-03",
"qwen-image-max",
"qwen-image-max-2025-12-30",
"qwen-image-plus",
"qwen-image-plus-2026-01-09",
"qwen-image",
}
_VALID_MODELS_EDIT = {
"qwen-image-2.0-pro",
"qwen-image-2.0-pro-2026-04-22",
"qwen-image-2.0-pro-2026-03-03",
"qwen-image-2.0",
"qwen-image-2.0-2026-03-03",
"qwen-image-edit-max",
"qwen-image-edit-max-2026-01-16",
"qwen-image-edit-plus",
"qwen-image-edit-plus-2025-12-15",
"qwen-image-edit-plus-2025-10-30",
"qwen-image-edit",
}
def _resolve_image_url(path_or_url: str) -> str:
"""Resolve an image path or URL to a usable string.
If the input is an HTTP/HTTPS URL, return it as-is.
If the input is a local file path, read the file and return
a base64 data URL.
Args:
path_or_url: HTTP/HTTPS URL or local file path.
Returns:
str: A URL (original URL or base64 data URL).
Raises:
FileNotFoundError: If the local file does not exist.
ValueError: If the file format is not supported.
"""
if path_or_url.startswith(("http://", "https://")):
return path_or_url
path_obj = Path(path_or_url)
if not path_obj.exists():
raise FileNotFoundError(
f"Image file not found: {path_or_url}",
)
if not path_obj.is_file():
raise ValueError(f"Not a file: {path_or_url}")
ext = path_obj.suffix.lower()
if ext not in _IMAGE_MIME_TYPES:
raise ValueError(
f"Unsupported image format: {ext}. "
f"Supported: {', '.join(_IMAGE_MIME_TYPES.keys())}",
)
mime_type = _IMAGE_MIME_TYPES[ext]
with open(path_obj, "rb") as f:
image_data = base64.b64encode(f.read()).decode("utf-8")
return f"data:{mime_type};base64,{image_data}"
def _extract_config(
tool_config: dict,
default_model: str,
) -> tuple[str, str, float, str]:
"""Extract api_key, endpoint, timeout and model from tool config.
Args:
tool_config: Tool configuration dict.
default_model: Fallback model name when not set in config.
Returns:
Tuple of (api_key, endpoint, timeout, model).
"""
api_key = tool_config.get("api_key", "")
endpoint = tool_config.get("endpoint", "")
if not endpoint or not endpoint.strip():
endpoint = _DEFAULT_ENDPOINT
timeout_raw = tool_config.get("timeout")
if timeout_raw is None or float(timeout_raw) <= 0:
timeout = _DEFAULT_TIMEOUT
else:
timeout = float(timeout_raw)
model = tool_config.get("model", "") or default_model
return api_key, endpoint, timeout, model
async def _download_image(
image_url: str,
save_dir: Path,
prefix: str,
timeout: float,
) -> Path:
"""Download image from URL and save to local directory.
Args:
image_url: Public URL of the image.
save_dir: Directory to save the image.
prefix: Filename prefix.
timeout: HTTP timeout in seconds.
Returns:
Path: Local path of the saved image.
Raises:
Exception: If download fails.
"""
save_dir.mkdir(parents=True, exist_ok=True)
timestamp = int(time.time() * 1000)
filename = f"{prefix}_{timestamp}.png"
image_path = save_dir / filename
async with httpx.AsyncClient(timeout=timeout) as client:
async with client.stream("GET", image_url) as response:
response.raise_for_status()
chunks = []
async for chunk in response.aiter_bytes(chunk_size=512 * 1024):
chunks.append(chunk)
await asyncio.to_thread(image_path.write_bytes, b"".join(chunks))
logger.info(f"Image saved to {image_path}")
return image_path
def _call_multimodal_conversation(
api_key: str,
endpoint: str,
model: str,
messages: list,
**kwargs,
):
"""Call DashScope MultiModalConversation with thread-safe setup.
Args:
api_key: DashScope API key.
endpoint: Base HTTP API URL.
model: Model name.
messages: Message list for the conversation.
**kwargs: Additional parameters passed to call().
Returns:
SDK response object.
"""
import dashscope
from dashscope import MultiModalConversation
with _DASHSCOPE_LOCK:
dashscope.base_http_api_url = endpoint
rsp = MultiModalConversation.call(
api_key=api_key,
model=model,
messages=messages,
result_format="message",
stream=False,
**kwargs,
)
return rsp
def _parse_image_urls(response) -> List[str]:
"""Extract image URLs from MultiModalConversation response.
Args:
response: SDK response object.
Returns:
List of image URL strings.
"""
urls = []
choices = getattr(
getattr(response, "output", None),
"choices",
None,
)
if not choices:
return urls
for choice in choices:
message = getattr(choice, "message", None)
if not message:
continue
content = getattr(message, "content", None)
if not content:
continue
for item in content:
if isinstance(item, dict):
url = item.get("image")
else:
url = getattr(item, "image", None)
if url:
urls.append(url)
return urls
async def generate_image_qwen(
prompt: str,
size: str = "2048*2048",
n: int = 1,
negative_prompt: str = "",
prompt_extend: bool = True,
) -> ToolChunk:
"""Generate images from a text prompt using Qwen-Image models.
Uses Alibaba Cloud's Qwen-Image models for high-quality image
generation. Supports complex text rendering, multi-style artwork,
and precise semantic adherence.
The model is selected via the tool's configuration settings.
Available models: qwen-image-2.0-pro (default), qwen-image-2.0,
qwen-image-max, qwen-image-plus, and dated snapshot versions.
Args:
prompt (str):
Text description of the image to generate.
Supports Chinese and English, up to 800 characters
(qwen-image-plus/max series) or longer (2.0 series).
size (str, optional):
Output image size in "width*height" format.
For qwen-image-2.0 series: total pixels between
512*512 and 2048*2048.
Recommended sizes: "2048*2048" (1:1, default),
"2688*1536" (16:9), "1536*2688" (9:16),
"2368*1728" (4:3).
For qwen-image-max/plus: use "1664*928" (16:9),
"1328*1328" (1:1), "928*1664" (9:16).
n (int, optional):
Number of images to generate (1-6 for 2.0 series,
fixed 1 for max/plus). Default: 1.
negative_prompt (str, optional):
Describe what to exclude from the image.
prompt_extend (bool, optional):
Enable prompt auto-optimization. Default: True.
Returns:
ToolChunk: Contains generated images and metadata.
"""
try:
tool_config = get_tool_config("generate_image_qwen")
if not tool_config:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
"Error: Tool not configured. "
"Please set your API key in the tool settings."
),
),
],
)
api_key, endpoint, timeout, model = _extract_config(
tool_config,
default_model="qwen-image-2.0-pro",
)
if not api_key:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
"Error: DashScope API key not configured. "
"Please set your API key in the tool settings."
),
),
],
)
if model not in _VALID_MODELS_GENERATE:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
f"Error: Invalid model '{model}'. "
f"Valid options: "
f"{', '.join(sorted(_VALID_MODELS_GENERATE))}"
),
),
],
)
if not 1 <= n <= 6:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
f"Error: Invalid n '{n}'. "
f"Must be between 1 and 6."
),
),
],
)
messages = [
{
"role": "user",
"content": [{"text": prompt}],
},
]
call_kwargs = {
"watermark": False,
"prompt_extend": prompt_extend,
"n": n,
}
if size:
call_kwargs["size"] = size
if negative_prompt:
call_kwargs["negative_prompt"] = negative_prompt
logger.info(
f"Generating image with Qwen-Image: "
f"model={model}, size={size}, n={n}",
)
rsp = await asyncio.to_thread(
_call_multimodal_conversation,
api_key=api_key,
endpoint=endpoint,
model=model,
messages=messages,
**call_kwargs,
)
if rsp.status_code != 200:
error_msg = (
f"DashScope API error: {rsp.status_code} - "
f"{rsp.code}: {rsp.message}"
)
logger.error(error_msg)
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=f"Error: {error_msg}",
),
],
)
image_urls = _parse_image_urls(rsp)
if not image_urls:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
"Error: No images returned from API. "
"Please try again."
),
),
],
)
logger.info(
f"Qwen-Image generated {len(image_urls)} image(s)",
)
save_dir = DEFAULT_MEDIA_DIR / "qwen_image"
content_blocks = []
saved_paths = []
for idx, img_url in enumerate(image_urls):
prefix = f"qwen_image_gen_{idx}"
try:
image_path = await _download_image(
img_url,
save_dir,
prefix,
timeout,
)
saved_paths.append(str(image_path))
content_blocks.append(
DataBlock(
source=URLSource(
url="file://" + str(image_path),
media_type=mimetypes.guess_type(
str(image_path),
)[0]
or "image/*",
),
),
)
except Exception as e:
logger.error(
f"Failed to download image {idx}: {e}",
)
content_blocks.append(
DataBlock(
source=URLSource(
url=img_url,
media_type=mimetypes.guess_type(img_url)[0]
or "image/*",
),
),
)
saved_paths.append(img_url)
content_blocks.append(
TextBlock(
type="text",
text=(
f"Generated {len(image_urls)} image(s) using "
f"Qwen-Image\n"
f"Model: {model}\n"
f"Prompt: {prompt}\n"
f"Size: {size}, Count: {n}\n"
f"Saved to: {', '.join(saved_paths)}"
),
),
)
return ToolChunk(state=ToolResultState.SUCCESS, content=content_blocks)
except Exception as e:
logger.error(
f"Qwen-Image generation failed: {e}",
exc_info=True,
)
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(f"Error: Image generation failed - {str(e)}"),
),
],
)
async def edit_image_qwen(
prompt: str,
reference_images: List[str],
size: str = "",
n: int = 1,
negative_prompt: str = "",
prompt_extend: bool = True,
) -> ToolChunk:
"""Edit or fuse images using Qwen-Image models.
Supports single-image editing (modify content, style transfer,
text rendering) and multi-image fusion (combine elements from
multiple images).
The model is selected via the tool's configuration settings.
Available models: qwen-image-2.0-pro (default), qwen-image-2.0,
qwen-image-edit-max, qwen-image-edit-plus, qwen-image-edit,
and dated snapshot versions.
Args:
prompt (str):
Description of the desired edit or fusion.
When multiple images are provided, use "图一" / "图二"
(or "image 1" / "image 2" in English) to refer to them.
reference_images (List[str]):
List of reference image URLs or local file paths
(.png/.jpg/.jpeg/.webp). At least 1 image required.
Each item can be:
- HTTP/HTTPS URL
- Local file path (auto-converted to base64)
size (str, optional):
Output image size in "width*height" format.
Leave empty to auto-detect based on input image.
For qwen-image-2.0 series: total pixels 512*512 to
2048*2048. Example: "1024*1024", "2048*2048".
n (int, optional):
Number of output images (1-6 for 2.0/edit-plus,
fixed 1 for edit-max/edit). Default: 1.
negative_prompt (str, optional):
Describe what to exclude from the output.
prompt_extend (bool, optional):
Enable prompt auto-optimization. Default: True.
Returns:
ToolChunk: Contains edited images and metadata.
"""
try:
if not reference_images:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
"Error: reference_images is required. "
"Please provide at least one image."
),
),
],
)
tool_config = get_tool_config("edit_image_qwen")
if not tool_config:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
"Error: Tool not configured. "
"Please set your API key in the tool settings."
),
),
],
)
api_key, endpoint, timeout, model = _extract_config(
tool_config,
default_model="qwen-image-2.0-pro",
)
if not api_key:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
"Error: DashScope API key not configured. "
"Please set your API key in the tool settings."
),
),
],
)
if model not in _VALID_MODELS_EDIT:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
f"Error: Invalid model '{model}'. "
f"Valid options: "
f"{', '.join(sorted(_VALID_MODELS_EDIT))}"
),
),
],
)
if not 1 <= n <= 6:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
f"Error: Invalid n '{n}'. "
f"Must be between 1 and 6."
),
),
],
)
# Build message content: images first, then prompt text
content = []
for img_input in reference_images:
try:
resolved = _resolve_image_url(img_input)
except (FileNotFoundError, ValueError) as e:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
f"Error: reference_images contains invalid "
f"entry '{img_input}' - {str(e)}"
),
),
],
)
content.append({"image": resolved})
content.append({"text": prompt})
messages = [{"role": "user", "content": content}]
call_kwargs = {
"watermark": False,
"prompt_extend": prompt_extend,
"n": n,
}
if size:
call_kwargs["size"] = size
if negative_prompt:
call_kwargs["negative_prompt"] = negative_prompt
logger.info(
f"Editing image with Qwen-Image: "
f"model={model}, "
f"reference_images={len(reference_images)}, n={n}",
)
rsp = await asyncio.to_thread(
_call_multimodal_conversation,
api_key=api_key,
endpoint=endpoint,
model=model,
messages=messages,
**call_kwargs,
)
if rsp.status_code != 200:
error_msg = (
f"DashScope API error: {rsp.status_code} - "
f"{rsp.code}: {rsp.message}"
)
logger.error(error_msg)
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=f"Error: {error_msg}",
),
],
)
image_urls = _parse_image_urls(rsp)
if not image_urls:
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=(
"Error: No images returned from API. "
"Please try again."
),
),
],
)
logger.info(
f"Qwen-Image edit produced {len(image_urls)} image(s)",
)
save_dir = DEFAULT_MEDIA_DIR / "qwen_image"
content_blocks = []
saved_paths = []
for idx, img_url in enumerate(image_urls):
prefix = f"qwen_image_edit_{idx}"
try:
image_path = await _download_image(
img_url,
save_dir,
prefix,
timeout,
)
saved_paths.append(str(image_path))
content_blocks.append(
DataBlock(
source=URLSource(
url="file://" + str(image_path),
media_type=mimetypes.guess_type(
str(image_path),
)[0]
or "image/*",
),
),
)
except Exception as e:
logger.error(
f"Failed to download image {idx}: {e}",
)
content_blocks.append(
DataBlock(
source=URLSource(
url=img_url,
media_type=mimetypes.guess_type(img_url)[0]
or "image/*",
),
),
)
saved_paths.append(img_url)
content_blocks.append(
TextBlock(
type="text",
text=(
f"Edited {len(image_urls)} image(s) using "
f"Qwen-Image\n"
f"Model: {model}\n"
f"Prompt: {prompt}\n"
f"Reference images: {len(reference_images)}\n"
f"Saved to: {', '.join(saved_paths)}"
),
),
)
return ToolChunk(state=ToolResultState.SUCCESS, content=content_blocks)
except Exception as e:
logger.error(
f"Qwen-Image edit failed: {e}",
exc_info=True,
)
return ToolChunk(
state=ToolResultState.ERROR,
content=[
TextBlock(
type="text",
text=f"Error: Image editing failed - {str(e)}",
),
],
)