#!/usr/bin/env python3 """Compare semantic inpainting quality on controlled, context-rich scenes.""" from __future__ import annotations import argparse import gc import json from pathlib import Path from typing import Any import numpy as np import torch from diffusers import ( DPMSolverMultistepScheduler, StableDiffusionInpaintPipeline, StableDiffusionPipeline, ) from PIL import Image, ImageDraw, ImageFilter from transformers import CLIPModel, CLIPProcessor SOURCE_CASES = ( { "name": "park-bench-cat", "source_prompt": ( "a detailed photograph of an empty wooden park bench centered in a leafy " "park, no people, no animals" ), "edit_prompt": ( "a tabby cat sitting naturally on the wooden park bench, detailed photography" ), "score_prompt": "a detailed photograph of a tabby cat", "mask": (148, 224, 366, 400), "shape": "ellipse", }, { "name": "kitchen-kettle", "source_prompt": ( "a detailed photograph of an empty kitchen countertop viewed straight on, " "warm daylight, no objects in the center" ), "edit_prompt": ( "a glossy red enamel kettle resting naturally on the kitchen countertop, " "detailed photography" ), "score_prompt": "a detailed photograph of a glossy red enamel kettle", "mask": (166, 236, 350, 414), "shape": "rounded_rectangle", }, { "name": "garden-greenhouse", "source_prompt": ( "a realistic moonlit garden with an empty grassy clearing in the center, " "lush plants around the clearing" ), "edit_prompt": ( "a tiny glass greenhouse glowing warmly in the moonlit garden clearing, " "detailed photography" ), "score_prompt": "a detailed photograph of a tiny glass greenhouse", "mask": ((170, 400), (147, 245), (193, 161), (321, 154), (372, 246), (350, 404)), "shape": "polygon", }, { "name": "street-bicycle", "source_prompt": ( "a detailed photograph of a quiet city street with an empty road in the " "foreground, late afternoon" ), "edit_prompt": ( "a bright red bicycle standing naturally on the city street, detailed photography" ), "score_prompt": "a detailed photograph of a bright red bicycle", "mask": ((146, 378), (206, 318), (288, 385), (366, 327)), "shape": "brush", "width": 92, }, { "name": "living-room-dog", "source_prompt": ( "a detailed photograph of a cozy living room with an empty rug centered on the " "floor, soft window light" ), "edit_prompt": ( "a small corgi sitting naturally on the living room rug, detailed photography" ), "score_prompt": "a detailed photograph of a small corgi dog", "mask": (156, 252, 360, 450), "shape": "ellipse", }, { "name": "lake-swan", "source_prompt": ( "a detailed photograph of a calm lake with empty water near the foreground, " "mountains in the distance" ), "edit_prompt": ( "a white swan floating naturally on the calm lake water, detailed photography" ), "score_prompt": "a detailed photograph of a white swan", "mask": ((151, 367), (208, 314), (272, 382), (355, 328)), "shape": "brush", "width": 96, }, ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--base_model", required=True) parser.add_argument("--base_revision") parser.add_argument("--baseline_model", required=True) parser.add_argument("--baseline_revision") parser.add_argument("--teacher_model", required=True) parser.add_argument("--teacher_revision") parser.add_argument("--teacher_variant") parser.add_argument("--candidate_model", required=True) parser.add_argument("--clip_model", default="openai/clip-vit-base-patch32") parser.add_argument("--clip_revision") parser.add_argument("--steps", type=int, default=30) parser.add_argument("--guidance_scale", type=float, default=7.5) parser.add_argument("--mask_crop_padding", type=int, default=0) parser.add_argument("--seed", type=int, default=20260811) parser.add_argument("--output_dir", type=Path, required=True) return parser.parse_args() def _scheduler(pipeline: Any) -> DPMSolverMultistepScheduler: return DPMSolverMultistepScheduler.from_config( pipeline.scheduler.config, algorithm_type="dpmsolver++", ) def _release_cuda(model: Any) -> None: if isinstance(model, torch.nn.Module): model.to("cpu") else: for component in getattr(model, "components", {}).values(): if isinstance(component, torch.nn.Module): component.to("cpu") del model gc.collect() torch.cuda.empty_cache() def _mask(case: dict[str, Any]) -> Image.Image: mask = Image.new("L", (512, 512), 0) draw = ImageDraw.Draw(mask) if case["shape"] == "ellipse": draw.ellipse(case["mask"], fill=255) elif case["shape"] == "rounded_rectangle": draw.rounded_rectangle(case["mask"], radius=28, fill=255) elif case["shape"] == "polygon": draw.polygon(case["mask"], fill=255) elif case["shape"] == "brush": points = case["mask"] width = int(case["width"]) draw.line(points, fill=255, width=width, joint="curve") radius = width // 2 for x, y in points: draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=255) else: raise ValueError(f"Unsupported semantic mask shape: {case['shape']}") return mask def _feather_inside(mask: Image.Image, radius: int = 6) -> Image.Image: binary = mask.convert("L").point(lambda value: 255 if value >= 128 else 0) softened = binary.filter(ImageFilter.GaussianBlur(radius=radius)) return Image.composite(softened, Image.new("L", binary.size, 0), binary) def _composite(generated: Image.Image, source: Image.Image, mask: Image.Image) -> Image.Image: return Image.composite(generated.convert("RGB"), source, _feather_inside(mask)) def _square_crop_bounds(mask: Image.Image, padding: int) -> tuple[int, int, int, int]: bounds = mask.convert("L").getbbox() if bounds is None: raise ValueError("Cannot crop around an empty mask") left, top, right, bottom = bounds side = min( max(mask.size), max(right - left, bottom - top) + max(0, padding) * 2, ) center_x = (left + right) / 2 center_y = (top + bottom) / 2 crop_left = round(center_x - side / 2) crop_top = round(center_y - side / 2) crop_left = min(max(0, crop_left), mask.width - side) crop_top = min(max(0, crop_top), mask.height - side) return crop_left, crop_top, crop_left + side, crop_top + side def _generate_sources(args: argparse.Namespace, cases: list[dict[str, Any]]) -> None: kwargs = {"revision": args.base_revision} if args.base_revision else {} pipeline = StableDiffusionPipeline.from_pretrained( args.base_model, torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False, **kwargs, ).to("cuda") pipeline.scheduler = _scheduler(pipeline) negative = "people, animals, object in the center, blurry, distorted, low detail" for index, case in enumerate(cases): generator = torch.Generator(device="cuda").manual_seed(args.seed + index) response = pipeline( prompt=case["source_prompt"], negative_prompt=negative, num_inference_steps=args.steps, guidance_scale=args.guidance_scale, width=512, height=512, generator=generator, ) case["source"] = response.images[0].convert("RGB") case["mask_image"] = _mask(case) case["source"].save(args.output_dir / f"source-{index:02d}-{case['name']}.png") case["mask_image"].save(args.output_dir / f"mask-{index:02d}-{case['name']}.png") _release_cuda(pipeline) def _run_inpainting_model( *, model_name: str, revision: str | None, variant: str | None, cases: list[dict[str, Any]], args: argparse.Namespace, output_dir: Path, ) -> list[Image.Image]: kwargs = {"revision": revision} if revision else {} if variant: kwargs.update({"variant": variant, "use_safetensors": True}) pipeline = StableDiffusionInpaintPipeline.from_pretrained( model_name, torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False, **kwargs, ).to("cuda") pipeline.scheduler = _scheduler(pipeline) results = [] for index, case in enumerate(cases): source = case["source"] mask_image = case["mask_image"] crop_bounds = None pipeline_image = source pipeline_mask = mask_image if args.mask_crop_padding > 0: crop_bounds = _square_crop_bounds(mask_image, args.mask_crop_padding) pipeline_image = source.crop(crop_bounds).resize( (512, 512), Image.Resampling.LANCZOS ) pipeline_mask = mask_image.crop(crop_bounds).resize( (512, 512), Image.Resampling.NEAREST ) generator = torch.Generator(device="cuda").manual_seed(args.seed + 10_000 + index) response = pipeline( prompt=case["edit_prompt"], negative_prompt="black patch, blurry, distorted, low detail", image=pipeline_image, mask_image=pipeline_mask, num_inference_steps=args.steps, guidance_scale=args.guidance_scale, width=512, height=512, generator=generator, ) generated = response.images[0] if crop_bounds is not None: generated = generated.resize( (crop_bounds[2] - crop_bounds[0], crop_bounds[3] - crop_bounds[1]), Image.Resampling.LANCZOS, ) full_generated = source.copy() full_generated.paste(generated, crop_bounds[:2]) generated = full_generated result = _composite(generated, source, mask_image) result.save(output_dir / f"{index:02d}-{case['name']}.png") results.append(result) _release_cuda(pipeline) return results def _crop_around_mask(image: Image.Image, mask: Image.Image, padding: int = 48) -> Image.Image: bounds = mask.getbbox() if bounds is None: return image left, top, right, bottom = bounds return image.crop( ( max(0, left - padding), max(0, top - padding), min(image.width, right + padding), min(image.height, bottom + padding), ) ) def _image_metrics( image: Image.Image, source: Image.Image, mask: Image.Image, ) -> dict[str, float | int]: image_array = np.asarray(image.convert("RGB"), dtype=np.int16) source_array = np.asarray(source.convert("RGB"), dtype=np.int16) selected = np.asarray(mask.convert("L")) >= 128 changed = np.any(image_array != source_array, axis=2) black = np.all(image_array <= 8, axis=2) absolute_change = np.abs(image_array - source_array).mean(axis=2) / 255.0 eroded = np.asarray( mask.convert("L").filter(ImageFilter.MinFilter(size=17)) ) >= 128 inner_boundary = selected & ~eroded return { "mask_area_fraction": float(selected.mean()), "masked_change_fraction": float(changed[selected].mean()), "masked_mean_absolute_change": float(absolute_change[selected].mean()), "boundary_mean_absolute_change": float( absolute_change[inner_boundary].mean() ), "masked_black_fraction": float(black[selected].mean()), "outside_changed_pixels": int(changed[~selected].sum()), } def _clip_scores( *, model_name: str, revision: str | None, cases: list[dict[str, Any]], outputs: dict[str, list[Image.Image]], ) -> tuple[dict[str, list[float]], dict[str, list[float]]]: kwargs = {"revision": revision} if revision else {} processor = CLIPProcessor.from_pretrained(model_name, **kwargs) model = CLIPModel.from_pretrained(model_name, **kwargs).to("cuda") scores: dict[str, list[float]] = {} for label, images in outputs.items(): label_scores = [] for case, image in zip(cases, images): inputs = processor( text=[case["score_prompt"]], images=[_crop_around_mask(image, case["mask_image"])], return_tensors="pt", padding=True, ).to("cuda") with torch.inference_mode(): vision = model.get_image_features(pixel_values=inputs["pixel_values"]) text = model.get_text_features( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], ) vision = vision / vision.norm(dim=-1, keepdim=True) text = text / text.norm(dim=-1, keepdim=True) label_scores.append(float((vision @ text.T).item())) scores[label] = label_scores teacher_scores = {label: [] for label in outputs} for index, case in enumerate(cases): labels = list(outputs) crops = [ _crop_around_mask(outputs[label][index], case["mask_image"]) for label in labels ] inputs = processor(images=crops, return_tensors="pt").to("cuda") with torch.inference_mode(): features = model.get_image_features(pixel_values=inputs["pixel_values"]) features = features / features.norm(dim=-1, keepdim=True) teacher_index = labels.index("teacher") similarities = features @ features[teacher_index] for label, similarity in zip(labels, similarities): teacher_scores[label].append(float(similarity.item())) _release_cuda(model) return scores, teacher_scores def _make_sheet( cases: list[dict[str, Any]], outputs: dict[str, list[Image.Image]], destination: Path, ) -> None: labels = ["source + mask", "current", "teacher", "candidate"] cell = 512 label_height = 34 sheet = Image.new("RGB", (cell * len(labels), (cell + label_height) * len(cases)), "white") draw = ImageDraw.Draw(sheet) for index, case in enumerate(cases): row_y = index * (cell + label_height) mask_overlay = Image.new("RGB", case["source"].size, (255, 255, 255)) source_mask = Image.blend( case["source"], Image.composite(mask_overlay, case["source"], case["mask_image"]), 0.55, ) images = [source_mask, outputs["baseline"][index], outputs["teacher"][index], outputs["candidate"][index]] for column, (label, image) in enumerate(zip(labels, images)): sheet.paste(image, (column * cell, row_y + label_height)) draw.text((column * cell + 8, row_y + 9), label, fill="black") draw.text((cell + 88, row_y + 9), f"{case['name']}: {case['edit_prompt'][:52]}", fill="black") sheet.save(destination) def main() -> None: args = parse_args() args.output_dir.mkdir(parents=True, exist_ok=False) cases = [dict(case) for case in SOURCE_CASES] _generate_sources(args, cases) outputs: dict[str, list[Image.Image]] = {} model_specs = { "baseline": (args.baseline_model, args.baseline_revision, None), "teacher": ( args.teacher_model, args.teacher_revision, args.teacher_variant, ), "candidate": (args.candidate_model, None, None), } for label, (model_name, revision, variant) in model_specs.items(): output_dir = args.output_dir / label output_dir.mkdir() outputs[label] = _run_inpainting_model( model_name=model_name, revision=revision, variant=variant, cases=cases, args=args, output_dir=output_dir, ) clip_scores, teacher_image_scores = _clip_scores( model_name=args.clip_model, revision=args.clip_revision, cases=cases, outputs=outputs, ) records = [] for index, case in enumerate(cases): record: dict[str, Any] = { "index": index, "name": case["name"], "source_prompt": case["source_prompt"], "edit_prompt": case["edit_prompt"], "score_prompt": case["score_prompt"], } for label, images in outputs.items(): record[label] = { **_image_metrics(images[index], case["source"], case["mask_image"]), "clip_similarity": clip_scores[label][index], "clip_similarity_to_teacher": teacher_image_scores[label][index], } records.append(record) summary = { "base_model": args.base_model, "base_revision": args.base_revision, "baseline_model": args.baseline_model, "baseline_revision": args.baseline_revision, "teacher_model": args.teacher_model, "teacher_revision": args.teacher_revision, "teacher_variant": args.teacher_variant, "candidate_model": args.candidate_model, "clip_model": args.clip_model, "clip_revision": args.clip_revision, "steps": args.steps, "guidance_scale": args.guidance_scale, "mask_crop_padding": args.mask_crop_padding, "seed": args.seed, "cases": records, "mean_clip_similarity": { label: float(np.mean(scores)) for label, scores in clip_scores.items() }, "mean_clip_similarity_to_teacher": { label: float(np.mean(scores)) for label, scores in teacher_image_scores.items() }, "candidate_clip_wins_over_baseline": int( sum( candidate > baseline for candidate, baseline in zip( clip_scores["candidate"], clip_scores["baseline"] ) ) ), } (args.output_dir / "metrics.json").write_text(json.dumps(summary, indent=2) + "\n") _make_sheet(cases, outputs, args.output_dir / "comparison.png") print(json.dumps(summary, indent=2)) if __name__ == "__main__": main()