384 lines
16 KiB
Python
Executable File
384 lines
16 KiB
Python
Executable File
import httpx
|
||
import uuid
|
||
import logging
|
||
|
||
from langchain_core.runnables import RunnableConfig
|
||
from minio import Minio
|
||
from langgraph.prebuilt import ToolRuntime
|
||
from src.core.config import settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
minio_client = Minio(settings.MINIO_URL, access_key=settings.MINIO_ACCESS, secret_key=settings.MINIO_SECRET, secure=settings.MINIO_SECURE)
|
||
|
||
from typing import List, Optional
|
||
from langchain_core.tools import tool
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@tool
|
||
async def generate_furniture(runtime: ToolRuntime, prompts: List[str] = None, num_images: Optional[int] = 12, ):
|
||
"""
|
||
生成家具设计线稿草图（sketch / line drawing）。
|
||
|
||
功能说明：
|
||
- 默认生成 12 张家具设计线稿。
|
||
- 智能处理 prompts 数量与生成数量不一致的情况：
|
||
- 如果只有一个 prompt → 用该 prompt 生成全部 12 张（不同随机变体）。
|
||
- 如果有多个 prompt → 自动均匀分配生成数量（尽量让每个 prompt 生成相同数量）。
|
||
- 生成过程会一张一张进行，适合用户实时查看。
|
||
|
||
参数说明：
|
||
- prompts (list[str]):
|
||
必须是列表，即使只有一个提示词也要用 ["你的提示词"] 格式。
|
||
提供详细的英文提示词，描述越详细越好。
|
||
- num_images (int, 可选): 要生成的图片总数量，默认 12 张，最大限制为 12 张。
|
||
|
||
返回值：
|
||
返回 image_urls 列表，系统会自动依次展示生成的图片。
|
||
"""
|
||
# ====================== 参数安全处理 ======================
|
||
if prompts is None or len(prompts) == 0:
|
||
return "Error: prompts 参数不能为空。请至少提供一个详细的英文提示词。"
|
||
|
||
if not isinstance(prompts, list):
|
||
prompts = [str(prompts)]
|
||
|
||
# 数量限制
|
||
if num_images is None or num_images < 1:
|
||
num_images = 1
|
||
elif num_images > 12:
|
||
num_images = 12
|
||
|
||
n_prompts = len(prompts)
|
||
|
||
logger.info(f"[generate_furniture] 开始生成 | prompts数量={n_prompts} | num_images={num_images}（默认12）")
|
||
|
||
# ====================== 均匀分配 prompts（核心逻辑） ======================
|
||
if n_prompts == 0:
|
||
return "Error: prompts 列表为空"
|
||
|
||
# 计算每个 prompt 应该生成的张数
|
||
base_count = num_images // n_prompts
|
||
remainder = num_images % n_prompts
|
||
|
||
images_per_prompt = [base_count] * n_prompts
|
||
for i in range(remainder):
|
||
images_per_prompt[i] += 1
|
||
|
||
# 构建实际使用的 prompt 列表
|
||
expanded_prompts: List[str] = []
|
||
for i, count in enumerate(images_per_prompt):
|
||
expanded_prompts.extend([prompts[i]] * count)
|
||
|
||
logger.info(f"[generate_furniture] 分配完成: {images_per_prompt} （每个prompt生成张数）")
|
||
|
||
# ====================== 生成图片 ======================
|
||
try:
|
||
bucket_name = "fida-public-bucket"
|
||
base_object_name = f"furniture/sketches/{uuid.uuid4()}"
|
||
image_urls = []
|
||
|
||
for i in range(num_images):
|
||
prompt = expanded_prompts[i]
|
||
object_name = f"{base_object_name}-{i:02d}.png"
|
||
|
||
image_url = await generate_or_edit_image(
|
||
prompt=prompt,
|
||
bucket_name=bucket_name,
|
||
object_name=object_name
|
||
)
|
||
image_urls.append(image_url)
|
||
|
||
logger.info(f"[generate_furniture] 已生成第 {i + 1}/{num_images} 张")
|
||
|
||
logger.info(f"[generate_furniture] 成功生成 {len(image_urls)} 张图片")
|
||
return image_urls
|
||
|
||
except Exception as e:
|
||
logger.error(f"generate_furniture 执行异常: {e}", exc_info=True)
|
||
return f"generate furniture error: {str(e)}"
|
||
|
||
|
||
@tool
|
||
async def edit_furniture(
|
||
runtime: ToolRuntime,
|
||
config: RunnableConfig,
|
||
input_image_paths: list[str] = None,
|
||
prompts: list[str] = None,
|
||
mode: str = "auto",
|
||
):
|
||
"""
|
||
使用先进的图像编辑模型对家具设计草图进行精准修改。
|
||
|
||
支持三种灵活模式（与 edit_quote_upload_furniture 保持一致）：
|
||
- one_to_one（默认，最常用）：多张图片 + 多个提示词，一一对应编辑
|
||
- one_to_many：1 张图片 + 多个提示词（同一张图片生成多个不同变体，例如不同风格/颜色）
|
||
- many_to_one：多张图片 + 1 个提示词（所有图片应用相同的修改）
|
||
|
||
参数说明：
|
||
- input_image_paths (list[str]): 输入图片的 MinIO 路径列表，长度建议 1~4
|
||
- prompts (list[str]): 修改提示词列表（必须是英文提示词）
|
||
- mode (str): "one_to_one", "one_to_many", "many_to_one", "auto"（默认自动判断）
|
||
|
||
使用要求（必须严格遵守）：
|
||
- input_image_paths 和 prompts 不能为空，长度必须在 1~4 之间。
|
||
- mode="auto" 时会根据列表长度智能判断模式：
|
||
- 1 张图片 + 多个 prompt → one_to_many
|
||
- 多个图片 + 1 个 prompt → many_to_one
|
||
- 图片数量 == prompt 数量 → one_to_one
|
||
- 编辑对象默认使用最近生成的图片（由 Supervisor 传入最新路径）。
|
||
|
||
示例调用：
|
||
|
||
1. one_to_one（一一对应，最常用）
|
||
input_image_paths = ["furniture/sketches/sofa_v1.png", "furniture/sketches/chair_v1.png"]
|
||
prompts = [
|
||
"Change the sofa to modern minimalist style with dark gray fabric.",
|
||
"Make the chair more Scandinavian with light wood and beige upholstery."
|
||
]
|
||
mode = "one_to_one"
|
||
|
||
2. one_to_many（同一张图片多个版本）
|
||
input_image_paths = ["furniture/sketches/sofa_latest.png"]
|
||
prompts = [
|
||
"Change to luxurious velvet with gold accents.",
|
||
"Change to industrial style with metal frame.",
|
||
"Change to soft pastel Nordic style."
|
||
]
|
||
mode = "one_to_many"
|
||
|
||
3. many_to_one（多张图片统一修改）
|
||
input_image_paths = ["furniture/sketches/sofa1.png", "furniture/sketches/chair1.png", "furniture/sketches/table1.png"]
|
||
prompts = ["Make all furniture more luxurious with velvet fabric and gold accents."]
|
||
mode = "many_to_one"
|
||
"""
|
||
try:
|
||
# ====================== 参数校验 ======================
|
||
if not input_image_paths or len(input_image_paths) < 1 or len(input_image_paths) > 4:
|
||
return f"参数错误：input_image_paths 必须提供，且长度需在 1 到 4 张之间。目前收到 {len(input_image_paths) if input_image_paths else 0} 张。"
|
||
|
||
if not prompts:
|
||
return "参数错误：prompts 不能为空，请至少提供一个修改提示词。"
|
||
|
||
if mode not in ["one_to_one", "one_to_many", "many_to_one", "auto"]:
|
||
return f"参数错误：mode 参数无效。可用值：one_to_one, one_to_many, many_to_one, auto。当前收到：{mode}"
|
||
|
||
# Auto 模式智能判断
|
||
if mode == "auto":
|
||
if len(input_image_paths) == 1 and len(prompts) > 1:
|
||
mode = "one_to_many"
|
||
elif len(prompts) == 1:
|
||
mode = "many_to_one"
|
||
elif len(input_image_paths) == len(prompts):
|
||
mode = "one_to_one"
|
||
else:
|
||
mode = "one_to_one" # 兜底
|
||
|
||
# 各模式严格校验
|
||
if mode == "one_to_many":
|
||
if len(input_image_paths) != 1:
|
||
return f"参数错误：one_to_many 模式只能传入 1 张图片，当前传入了 {len(input_image_paths)} 张。"
|
||
if len(prompts) < 1:
|
||
return "参数错误：one_to_many 模式下 prompts 至少需要 1 个。"
|
||
|
||
elif mode == "many_to_one":
|
||
if len(prompts) != 1:
|
||
return f"参数错误：many_to_one 模式下 prompts 必须只有 1 个，当前有 {len(prompts)} 个。"
|
||
|
||
elif mode == "one_to_one":
|
||
if len(prompts) != len(input_image_paths):
|
||
return (f"参数错误：one_to_one 模式下 input_image_paths 和 prompts 数量必须完全一致。\n"
|
||
f"当前图片 {len(input_image_paths)} 张，prompts {len(prompts)} 个。")
|
||
|
||
# ====================== 执行编辑 ======================
|
||
result = []
|
||
bucket_name = "fida-public-bucket"
|
||
|
||
if mode == "one_to_many":
|
||
# 同一张图片 + 多个 prompt
|
||
base_image = input_image_paths[0]
|
||
for i, prompt in enumerate(prompts):
|
||
object_name = f"furniture/sketches/{uuid.uuid4()}.png"
|
||
image_url = await generate_or_edit_image(
|
||
input_path=[base_image],
|
||
prompt=prompt,
|
||
bucket_name=bucket_name,
|
||
object_name=f"{object_name}-var{i}.png"
|
||
)
|
||
result.append(image_url)
|
||
|
||
elif mode == "many_to_one":
|
||
# 多张图片 + 1 个 prompt
|
||
current_prompt = prompts[0]
|
||
for i, image_path in enumerate(input_image_paths):
|
||
object_name = f"furniture/sketches/{uuid.uuid4()}.png"
|
||
image_url = await generate_or_edit_image(
|
||
input_path=[image_path],
|
||
prompt=current_prompt,
|
||
bucket_name=bucket_name,
|
||
object_name=f"{object_name}-{i}.png"
|
||
)
|
||
result.append(image_url)
|
||
|
||
else:
|
||
# one_to_one：一一对应
|
||
for i in range(len(input_image_paths)):
|
||
object_name = f"furniture/sketches/{uuid.uuid4()}.png"
|
||
image_url = await generate_or_edit_image(
|
||
input_path=[input_image_paths[i]],
|
||
prompt=prompts[i],
|
||
bucket_name=bucket_name,
|
||
object_name=f"{object_name}-{i}.png"
|
||
)
|
||
result.append(image_url)
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"edit_furniture 执行异常: {e}", exc_info=True)
|
||
return f"工具执行失败：{str(e)}，请检查参数后重试。"
|
||
|
||
|
||
@tool
|
||
async def edit_quote_upload_furniture(image_paths: list[str] = None, mode: str = "auto", prompts: list[str] = None, ):
|
||
"""
|
||
使用先进的图像编辑模型对家具图片进行精准批量修改。
|
||
|
||
支持四种模式：
|
||
- one_to_one（最常用）：多张图片 + 多个提示词，一一对应编辑
|
||
- one_to_many：多张图片 + 1个提示词（所有图片统一修改）
|
||
- many_to_one：1张图片 + 多个提示词（同一张图生成多个不同变体，例如不同颜色）
|
||
- many_to_many（新增）：多张图片 + 多个提示词，一一对应（多对多交叉编辑）
|
||
|
||
参数说明：
|
||
- image_paths (list[str]): MinIO 图片路径列表，长度建议 1~4
|
||
- prompts (list[str]): 详细英文提示词列表
|
||
- mode (str): "one_to_one", "one_to_many", "many_to_one", "many_to_many", "auto"（默认自动判断）
|
||
|
||
使用要求：
|
||
- image_paths 长度必须在 1~4 之间
|
||
- mode="auto" 时会根据长度智能判断
|
||
- many_to_many 模式下：image_paths 和 prompts 的长度必须完全相同
|
||
|
||
示例：
|
||
|
||
示例1：many_to_many（多对多，一一对应）
|
||
image_paths = ["sofa1.png", "chair1.png", "table1.png"]
|
||
prompts = [
|
||
"Change to bright yellow modern style.",
|
||
"Change to deep green luxury style.",
|
||
"Change to soft beige Scandinavian style."
|
||
]
|
||
mode = "many_to_many"
|
||
|
||
示例2：many_to_one（同一张图多个颜色版本）
|
||
image_paths = ["sofa_original.png"]
|
||
prompts = ["yellow version", "green version", "blue version", "black version"]
|
||
mode = "many_to_one"
|
||
"""
|
||
try:
|
||
# ====================== 参数校验（直接返回错误信息） ======================
|
||
if not image_paths or len(image_paths) < 1 or len(image_paths) > 4:
|
||
return f"参数错误：image_paths 必须提供，且长度需要在 1 到 4 张之间。目前收到 {len(image_paths) if image_paths else 0} 张。"
|
||
|
||
if not prompts:
|
||
return "参数错误：prompts 不能为空，请至少提供一个修改提示词。"
|
||
|
||
if mode not in ["one_to_one", "one_to_many", "many_to_one", "many_to_many", "auto"]:
|
||
return f"参数错误：mode 参数无效。可用值：one_to_one, one_to_many, many_to_one, many_to_many, auto。当前收到：{mode}"
|
||
|
||
# Auto 模式智能判断
|
||
if mode == "auto":
|
||
if len(image_paths) == 1 and len(prompts) > 1:
|
||
mode = "many_to_one"
|
||
elif len(prompts) == 1:
|
||
mode = "one_to_many"
|
||
elif len(image_paths) == len(prompts):
|
||
mode = "many_to_many" # 新增：数量相等时优先 many_to_many
|
||
else:
|
||
mode = "one_to_one"
|
||
|
||
# 各模式严格校验
|
||
if mode == "many_to_one":
|
||
if len(image_paths) != 1:
|
||
return f"参数错误：many_to_one 模式只能传入 1 张图片，当前传入了 {len(image_paths)} 张。"
|
||
if len(prompts) < 1:
|
||
return "参数错误：many_to_one 模式下 prompts 至少需要 1 个。"
|
||
|
||
elif mode == "one_to_many":
|
||
if len(prompts) != 1:
|
||
return f"参数错误：one_to_many 模式下 prompts 必须只有 1 个，当前有 {len(prompts)} 个。"
|
||
|
||
elif mode in ["one_to_one", "many_to_many"]:
|
||
if len(prompts) != len(image_paths):
|
||
return (f"参数错误：{mode} 模式下 image_paths 和 prompts 数量必须完全一致。\n"
|
||
f"当前 image_paths 有 {len(image_paths)} 张，prompts 有 {len(prompts)} 个。")
|
||
|
||
# ====================== 执行编辑 ======================
|
||
result = []
|
||
bucket_name = "fida-public-bucket"
|
||
|
||
if mode == "many_to_one":
|
||
# 同一张图片 + 多个 prompt
|
||
base_image = image_paths[0]
|
||
for i, prompt in enumerate(prompts):
|
||
object_name = f"furniture/sketches/{uuid.uuid4()}.png"
|
||
image_url = await generate_or_edit_image(
|
||
input_path=[base_image],
|
||
prompt=prompt,
|
||
bucket_name=bucket_name,
|
||
object_name=f"{object_name}-var{i}.png"
|
||
)
|
||
result.append(image_url)
|
||
|
||
else:
|
||
# one_to_one、many_to_many、one_to_many 统一处理
|
||
for i in range(len(image_paths)):
|
||
# 根据模式决定当前使用的 prompt
|
||
if mode == "one_to_many":
|
||
current_prompt = prompts[0]
|
||
else:
|
||
current_prompt = prompts[i] # one_to_one 和 many_to_many 都用对应位置的 prompt
|
||
|
||
object_name = f"furniture/sketches/{uuid.uuid4()}.png"
|
||
image_url = await generate_or_edit_image(
|
||
input_path=[image_paths[i]],
|
||
prompt=current_prompt,
|
||
bucket_name=bucket_name,
|
||
object_name=f"{object_name}-{i}.png"
|
||
)
|
||
result.append(image_url)
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"edit_quote_upload_furniture 执行异常: {e}", exc_info=True)
|
||
return f"工具执行失败：{str(e)}，请检查参数后重试。"
|
||
|
||
|
||
async def generate_or_edit_image(input_path=None, bucket_name="fida-public-bucket",
|
||
object_name=f"furniture/sketches/{uuid.uuid4()}.png",
|
||
prompt="Generate a modern minimalist dining chair made of light "
|
||
"oak wood and white leather, with slim metal legs, photographed "
|
||
"in a bright Scandinavian living room with natural sunlight, high detail, "
|
||
"8k resolution."):
|
||
if input_path is None:
|
||
input_path = []
|
||
request_data = {
|
||
"input_image_paths": input_path,
|
||
"prompt": prompt,
|
||
"bucket_name": bucket_name,
|
||
"object_name": object_name,
|
||
"width": 1024,
|
||
"height": 1024
|
||
}
|
||
async with httpx.AsyncClient(timeout=120) as client:
|
||
resp = await client.post(
|
||
f"http://{settings.FLUX2_GEN_IMG_MODEL_URL}/predict",
|
||
json=request_data,
|
||
)
|
||
result = resp.json()
|
||
image_url = result.get("output_path", None)
|
||
return image_url
|