feat: add outfit generation core

src/core/ai_outfit.py

@@ -0,0 +1,131 @@
+import logging
+import re
+from io import BytesIO
+from pathlib import Path
+
+from PIL import Image, ImageOps
+
+from core.models import OutfitResult, OutfitTask
+from services.ai_image_service import ImageApiClient
+
+logger = logging.getLogger(__name__)
+
+QUALITY_SMALL = 75
+QUALITY_BALANCED = 85
+QUALITY_HIGH = 92
+QUALITY_PRESETS = {
+    "small": QUALITY_SMALL,
+    "balanced": QUALITY_BALANCED,
+    "high": QUALITY_HIGH,
+    "小文件": QUALITY_SMALL,
+    "均衡": QUALITY_BALANCED,
+    "高清": QUALITY_HIGH,
+}
+
+MAX_JPG_BYTES = 2 * 1024 * 1024
+_INVALID_FILENAME_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')
+
+
+def render_prompt(template, task):
+    """Render an outfit prompt for one task."""
+    return str(template).replace("{title}", task.title).replace("{product_id}", task.product_id)
+
+
+def safe_product_filename(product_id):
+    """Return a filesystem-safe base filename without changing Excel data."""
+    name = _INVALID_FILENAME_CHARS.sub("_", str(product_id))
+    name = name.strip().strip(".")
+    return name or "outfit"
+
+
+def make_outfit_output_path(output_dir, product_id):
+    """Return output_dir/<product_id>.jpg without overwriting existing files."""
+    target_dir = Path(output_dir)
+    stem = safe_product_filename(product_id)
+    candidate = target_dir / (stem + ".jpg")
+    counter = 1
+    while candidate.exists():
+        candidate = target_dir / ("{}_{}.jpg".format(stem, counter))
+        counter += 1
+    return candidate
+
+
+def save_jpg_under_limit(image_bytes, output_path, quality=QUALITY_BALANCED, max_bytes=MAX_JPG_BYTES):
+    """Save image bytes as 1:1 JPG, reducing quality/size until under limit."""
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with Image.open(BytesIO(image_bytes)) as opened:
+        img = ImageOps.exif_transpose(opened).convert("RGB")
+        img = _crop_square(img)
+
+    quality = _coerce_quality(quality)
+    min_quality = 55
+    while True:
+        _save_jpeg(img, output_path, quality)
+        if output_path.stat().st_size <= max_bytes:
+            return output_path
+        if quality > min_quality:
+            quality = max(min_quality, quality - 7)
+            continue
+        if img.width <= 512 or img.height <= 512:
+            return output_path
+        new_size = max(512, int(img.width * 0.85))
+        img = img.resize((new_size, new_size), Image.LANCZOS)
+        quality = _coerce_quality(quality)
+
+
+def generate_outfit_image(
+    task,
+    prompt_template,
+    output_dir,
+    model_config,
+    quality=QUALITY_BALANCED,
+    resolution="1K",
+    api_client=None,
+):
+    """Generate one outfit image and return OutfitResult. Never raises."""
+    if not isinstance(task, OutfitTask):
+        raise TypeError("task must be OutfitTask")
+
+    try:
+        prompt = render_prompt(prompt_template, task)
+        client = api_client or ImageApiClient(model_config)
+        image_bytes = client.generate(prompt, task.garment_path, resolution=resolution)
+        output_path = make_outfit_output_path(output_dir, task.product_id)
+        save_jpg_under_limit(image_bytes, output_path, quality=quality)
+        logger.info("Generated outfit row %s -> %s", task.row_index, output_path)
+        return OutfitResult(
+            task=task,
+            success=True,
+            output_path=str(output_path),
+            attempts=1,
+        )
+    except Exception as exc:
+        logger.exception("Outfit generation failed for row %s", task.row_index)
+        return OutfitResult(
+            task=task,
+            success=False,
+            error=str(exc),
+            attempts=1,
+        )
+
+
+def _coerce_quality(quality):
+    if isinstance(quality, str):
+        return QUALITY_PRESETS.get(quality, QUALITY_BALANCED)
+    try:
+        return max(1, min(95, int(quality)))
+    except (TypeError, ValueError):
+        return QUALITY_BALANCED
+
+
+def _crop_square(img):
+    width, height = img.size
+    side = min(width, height)
+    left = (width - side) // 2
+    top = (height - side) // 2
+    return img.crop((left, top, left + side, top + side))
+
+
+def _save_jpeg(img, output_path, quality):
+    img.save(str(output_path), format="JPEG", quality=quality, optimize=True)