"""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)