core.py

"""CPU-side dataset export and time-budget orchestration; independently testable."""
import io
import tempfile
import time
import zipfile
from dataclasses import dataclass, field
from pathlib import Path

from PIL import Image, ImageChops, ImageOps

from views import BY_KEY, DEFAULT_KEYS, SIZES, make_prompt, validate_keys


class BudgetExhausted(Exception):
    pass


@dataclass
class BatchResult:
    images: dict = field(default_factory=dict)
    timings: list = field(default_factory=list)
    stop_reason: str = ""


def rgb_on_white(image):
    """Return an independent RGB image, compositing any transparency onto white."""
    rgba = image.convert("RGBA")
    white = Image.new("RGBA", rgba.size, "white")
    white.alpha_composite(rgba)
    return white.convert("RGB")


def prepare_reference(image, max_edge):
    image = ImageOps.exif_transpose(image)
    result = rgb_on_white(image)
    result.thumbnail((max_edge, max_edge), Image.Resampling.LANCZOS)
    return result


def conditioning_reference(reference, key, width, height):
    """Crop close views from an upright full-body reference on a white background.

    Fractions are measured from the detected foreground's head to feet. Crops are
    derived independently from the original reference, with a small head margin,
    and fitted proportionally onto a white canvas at the requested output size.
    Full-body/pose views preserve the original reference object unchanged.
    """
    fraction = {"cowboy_front": 0.72, "cowboy_back": 0.72,
                "upper_front": 0.45, "face_front": 0.20}.get(key)
    if fraction is None:
        return reference
    if width <= 0 or height <= 0:
        raise ValueError("Conditioning dimensions must be positive")
    rgb = rgb_on_white(reference)
    difference = ImageChops.difference(rgb, Image.new("RGB", rgb.size, "white"))
    red, green, blue = difference.split()
    distance = ImageChops.lighter(ImageChops.lighter(red, green), blue)
    mask = distance.point(lambda value: 255 if value > 18 else 0)
    bbox = mask.getbbox() or (0, 0, rgb.width, rgb.height)
    subject_height = bbox[3] - bbox[1]
    keep_height = max(1, round(subject_height * fraction))
    top_padding = max(1, round(keep_height * 0.06))
    top = max(0, bbox[1] - top_padding)
    bottom = min(rgb.height, bbox[1] + keep_height)

    # Use this band's foreground width so the full body's arms/legs do not force
    # excess whitespace around a face crop. Preserve a small horizontal margin.
    band_bbox = mask.crop((0, top, rgb.width, bottom)).getbbox()
    left, right = (band_bbox[0], band_bbox[2]) if band_bbox else (bbox[0], bbox[2])
    side_padding = max(1, round((right - left) * 0.06))
    crop = rgb.crop((max(0, left - side_padding), top,
                    min(rgb.width, right + side_padding), bottom))
    fitted = ImageOps.contain(crop, (width, height), Image.Resampling.LANCZOS)
    canvas = Image.new("RGB", (width, height), "white")
    canvas.paste(fitted, ((width - fitted.width) // 2, (height - fitted.height) // 2))
    return canvas


def requested_seconds(keys, maximum=110):
    """An execution budget, NOT a benchmark-based estimate or a guarantee."""
    # Even a one-view smoke test needs startup and decode time in addition to
    # the denoising budget. Keep a practical floor without claiming a speed.
    return min(maximum, max(45, 15 + 10 * len(validate_keys(keys))))


def run_batch(generate_one, reference, keys, guide_images, resolution, seed, extra,
              angles, budget_seconds, clock=time.monotonic, crop_close_views=False):
    keys = validate_keys(keys)
    result = BatchResult()
    started = clock()
    deadline = started + budget_seconds
    longest = 0.0

    def check_budget():
        # Called at denoising step boundaries. Leave room for decode and result transfer.
        if clock() >= deadline - 6:
            raise BudgetExhausted()

    for key in keys:
        remaining = deadline - clock()
        if remaining <= max(8.0, longest * 1.4 + 6):
            result.stop_reason = "時間枠に達する前に終了しました。完成した画像を保存できます。"
            break
        view = BY_KEY[key]
        width, height = SIZES[resolution][view.shape]
        inputs = [conditioning_reference(reference, key, width, height)
                  if crop_close_views else reference]
        if guide_images is not None:
            inputs.append(guide_images[key])
        # Stable per-view seed, unaffected by selecting or reordering other views.
        view_seed = (int(seed) + DEFAULT_KEYS.index(key)) % (2**31)
        t0 = clock()
        try:
            output = generate_one(
                image=inputs, prompt=make_prompt(view, extra, guide_images is not None, angles),
                width=width, height=height, seed=view_seed, check_budget=check_budget,
            )
        except BudgetExhausted:
            result.stop_reason = "時間枠に達する前に終了しました。完成した画像を保存できます。"
            break
        except Exception as exc:
            # A late view failure must not discard images already completed.
            print(f"[generation] {key} failed: {type(exc).__name__}: {exc}", flush=True)
            result.stop_reason = f"「{view.label}」の生成に失敗しました。完成済み画像は保存できます。"
            break
        elapsed = clock() - t0
        longest = max(longest, elapsed)
        result.images[key] = output
        result.timings.append({"key": key, "seconds": round(elapsed, 3), "seed": view_seed,
                               "width": width, "height": height, "references": len(inputs)})
        print(f"[generation] {key}: {elapsed:.2f}s", flush=True)
    return result


def build_zip(images, selected_keys, output_dir=None):
    """Export PNGs only; optionally write inside a caller-managed directory."""
    keys = validate_keys(selected_keys)
    if not any(key in images for key in keys):
        return None
    folder = Path(output_dir) if output_dir is not None else Path(tempfile.mkdtemp(prefix="character-views-"))
    folder.mkdir(parents=True, exist_ok=True)
    target = folder / "character_views.zip"
    with zipfile.ZipFile(target, "w", zipfile.ZIP_STORED) as archive:
        for key in keys:
            if key not in images:
                continue
            buffer = io.BytesIO()
            rgb_on_white(images[key]).save(buffer, format="PNG")
            index = DEFAULT_KEYS.index(key) + 1
            archive.writestr(f"dataset/{index:02d}_{key}.png", buffer.getvalue())
    return str(target)