Document and package context-aware inpainting v2

Adds selected training provenance, release metrics, real inpainting examples, runtime recommendations, and current training/evaluation sources.

.gitattributes

@@ -35,3 +35,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 examples/result-cat.png filter=lfs diff=lfs merge=lfs -text
 examples/source-greenhouse.png filter=lfs diff=lfs merge=lfs -text
+examples/bicycle-result.png filter=lfs diff=lfs merge=lfs -text
+examples/bicycle-source.png filter=lfs diff=lfs merge=lfs -text
+examples/cat-result.png filter=lfs diff=lfs merge=lfs -text
+examples/cat-source.png filter=lfs diff=lfs merge=lfs -text
+examples/kettle-result.png filter=lfs diff=lfs merge=lfs -text
+examples/kettle-source.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -12,102 +12,140 @@ tags:
 
 # Clover Image Tiny Inpaint 🍀
 
-An inpainting adaptation of Clover Image Tiny for 512×512 local generation and
-on-device Core ML deployment. White mask pixels are regenerated; black pixels
-are preserved. The intended target is an SD 1.4-class architecture, rather than
-a device-specific model.
+Clover Image Tiny Inpaint is a compact, context-aware inpainting model for the
+SD 1.4-class 512×512 architecture. Paint the region to replace in white; black
+pixels are preserved. This v2 checkpoint was selected from bounded
+teacher-distillation sweeps for semantic prompt following, reconstruction, and
+clean mask-edge blending.
 
-The model uses a 9-channel U-Net input:
+The model uses the standard nine-channel inpainting contract:
 
 ```text
 [noisy latent (4), mask (1), masked-image latent (4)]
 ```
 
-The base text encoder, VAE, scheduler, safety checker, and tokenizer remain
-compatible with Clover Image Tiny. The inpainting export additionally includes
-the VAE encoder needed to prepare the masked-image latent on device.
+The text encoder, VAE, scheduler, safety checker, and tokenizer remain
+compatible with Clover Image Tiny. A VAE encoder is additionally required to
+prepare the masked-image latent.
 
 ## Diffusers example
 
 ```python
+import torch
 from diffusers import AutoPipelineForInpainting, DPMSolverMultistepScheduler
 from diffusers.utils import load_image
 
 pipe = AutoPipelineForInpainting.from_pretrained(
     "neonforestmist/Clover-Image-Tiny-Inpaint",
-    torch_dtype="auto",
-)
+    torch_dtype=torch.float16,
+).to("cuda")
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-image = pipe(
-    prompt="a tiny glass greenhouse glowing in a moonlit garden",
-    image=load_image("input.png"),
+
+result = pipe(
+    prompt="a glossy red enamel kettle resting naturally on the countertop",
+    negative_prompt="blurry, distorted, low detail",
+    image=load_image("source.png"),
     mask_image=load_image("mask.png"),
     num_inference_steps=20,
-    padding_mask_crop=64,
+    guidance_scale=6.0,
+    padding_mask_crop=96,
 ).images[0]
-image.save("clover-inpaint.png")
+result.save("clover-inpaint.png")
 ```
 
-The mask is a grayscale image: white means “regenerate” and black means
-“preserve.” A few useful inpainting prompts are:
+Recommended interactive defaults are DPM-Solver++, 20 steps, CFG 6.0, and
+`padding_mask_crop=96`. Keep runs at or below 50 steps. For production UI,
+composite the generated image through the exact binary mask so that every
+unmasked source pixel remains unchanged.
+
+## Inpainting examples
 
-- `replace the masked area with a tiny glass greenhouse glowing at night`
-- `remove the person from the masked area and continue the background naturally`
-- `add a red enamel kettle on the masked countertop`
+All examples below use the selected v2 checkpoint, CFG 6.0, 30 evaluation
+steps, and a 96-pixel context crop. White is the region regenerated by Clover.
 
-## Inpainting example
+### Context-aware object insertion
 
-The example below uses DPM-Solver++, 20 steps, CFG 7.5, and seed 1337. The
-prompt is `a realistic orange cat sitting in the doorway, detailed photography`.
-Only the white doorway mask is replaced.
+Prompt: `a tabby cat sitting naturally on the wooden park bench`
 
 | Source | White mask | Result |
 |:---:|:---:|:---:|
-| ![Greenhouse source](examples/source-greenhouse.png) | ![White doorway mask](examples/mask-doorway.png) | ![Cat inpainted into the doorway](examples/result-cat.png) |
+| ![Park bench source](examples/cat-source.png) | ![Cat edit mask](examples/cat-mask.png) | ![Cat inpaint result](examples/cat-result.png) |
 
-For small masks, `padding_mask_crop=64` gives the edit more useful latent
-resolution before Diffusers maps it back into the source image. DPM-Solver++
-is the recommended scheduler for this release; PNDM can become unstable at
-higher step counts. Start at 20 steps and keep interactive runs at or below 50.
+Prompt: `a glossy red enamel kettle resting naturally on the countertop`
 
-For repeatable experiments, keep the input image, mask, seed, scheduler, and
-step count together. Inpainting is local editing: the unmasked region is
-provided as the masked-image conditioning and is also preserved by the native
-runtime compositor.
+| Source | White mask | Result |
+|:---:|:---:|:---:|
+| ![Kitchen source](examples/kettle-source.png) | ![Kettle edit mask](examples/kettle-mask.png) | ![Kettle inpaint result](examples/kettle-result.png) |
 
-## Core ML and SD 1.4-class deployment
+### Irregular mask replacement
 
-The companion Core ML resource bundle is converted for iOS 18 with a
-batch-one U-Net and chunked U-Net resources for the SD 1.4-class 512×512
-architecture. The Swift runtime performs classifier-free guidance as two
-serial passes to reduce peak memory. The
-bundled `VAEEncoder.mlmodelc` creates the masked-image latent locally, so the
-input image and mask do not leave the device.
+Prompt: `a bright red bicycle standing naturally on the city street`
 
-Conversion and the native iOS integration live in the source Clover repo:
+| Source | White mask | Result |
+|:---:|:---:|:---:|
+| ![Street source](examples/bicycle-source.png) | ![Bicycle edit mask](examples/bicycle-mask.png) | ![Bicycle inpaint result](examples/bicycle-result.png) |
 
-- [`coreml-tools/convert_inpaint.sh`](https://huggingface.co/neonforestmist/Clover-Image-Tiny/blob/main/coreml-tools/convert_inpaint.sh)
-- [`Clover-iOS`](https://huggingface.co/neonforestmist/Clover-Image-Tiny/tree/main/Clover-iOS)
-- [`training/README-INPAINTING.md`](https://huggingface.co/neonforestmist/Clover-Image-Tiny/blob/main/training/README-INPAINTING.md)
+## Evaluation
 
-## LoRA compatibility
+The release gate used a deterministic 12-image held-out reconstruction set and
+six context-rich semantic edits with ellipse, rounded-rectangle, polygon, and
+brush masks.
 
-Diffusers exposes LoRA loading for inpainting pipelines, but an adapter must be
-trained against this model's 9-channel U-Net. Existing Clover Image Tiny style
-LoRAs target the Regular 4-channel U-Net and are not interchangeable. The
-current Core ML export does not expose dynamic LoRA state; an inpainting-specific
-adapter can instead be fused and converted as a separate bundle.
+| Metric | Previous release | v2 |
+|---|---:|---:|
+| Held-out masked-region MAE (lower is better) | 0.3120 | **0.2517** |
+| Mean object-text CLIP similarity | 0.2726 | **0.2890** |
+| Mean CLIP image similarity to the SD inpainting teacher | 0.7651 | **0.8183** |
+| Black-collapse outputs | 0/6 | **0/6** |
+| Changed pixels outside the mask | 0 | **0** |
+
+Masked reconstruction MAE improved by 19.3%. Metrics are useful regression
+signals, not guarantees of photorealism; the visual contact sheet and per-case
+outputs were also reviewed before selection.
 
 ## Training provenance
 
-Training uses synthetic rectangle, ellipse, and brush masks over the pinned
-Apache-2.0 `prithivMLmods/Caption3o-Opt` image-caption dataset. The full job is
-launched by Modal under the `guccichungus69` workspace and stores its output in
-the `clover-image-tiny-inpaint-output` Volume before Core ML conversion.
+- Warm start: `neonforestmist/Clover-Image-Tiny-Inpaint` at revision
+  `1b6f8ae3db51900520369d5522c7dc7c2a97e21e`
+- Teacher: `stable-diffusion-v1-5/stable-diffusion-inpainting` at revision
+  `8a4288a76071f7280aedbdb3253bdb9e9d5d84bb`
+- Dataset: `prithivMLmods/Caption3o-Opt` at revision
+  `17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b`
+- Optimizer steps: 500, with cosine decay after warmup
+- Objective: teacher noise prediction plus ground-truth denoising, Min-SNR
+  weighting, and extra masked-region/boundary weighting
+- Masks: brush, multi-brush, rectangle, ellipse, polygon, multi-region, and
+  outpainting distributions
+- Caption dropout: 0.1
+- Training platform: Modal A10 in the `guccichungus69` workspace
+
+The longer run was rejected after its held-out and semantic gates regressed;
+the published checkpoint is the best-performing bounded sweep, not simply the
+last checkpoint.
+
+## Core ML
+
+The companion SD 1.4-class Core ML resources are published at
+[`neonforestmist/Clover-Image-Tiny-Inpaint-CoreML`](https://huggingface.co/neonforestmist/Clover-Image-Tiny-Inpaint-CoreML).
+The batch-one U-Net accepts `[1, 9, 64, 64]`; classifier-free guidance is run as
+two serial U-Net passes to reduce peak memory.
 
-## Citation
+## LoRA compatibility
+
+Diffusers can load a LoRA trained against this nine-channel inpainting U-Net.
+Regular Clover Image Tiny LoRAs target a four-channel U-Net and are not
+interchangeable. The Core ML package does not dynamically load inpainting
+LoRAs; fuse an inpainting-specific adapter before conversion if needed.
 
-If Clover Image Tiny Inpaint is useful in your work, please cite the release:
+## Limitations
+
+The model can still distort small, highly structured objects, text, hands, and
+faces. Very small masks may not provide enough latent resolution without the
+recommended context crop. Results depend on the source, mask, prompt, seed,
+and scheduler. This release inherits the limitations and license obligations
+of its base and teacher models.
+
+## Citation
 
 ```bibtex
 @software{lozadaperez2026cloverimagetinyinpaint,

inpainting-config.json

@@ -1,6 +1,63 @@
 {
+  "recipe_version": 2,
+  "base_model": "neonforestmist/Clover-Image-Tiny",
+  "base_revision": "63b0e9f6be9c00888ff464f342a9ef052bf76681",
+  "initial_inpaint_model": "neonforestmist/Clover-Image-Tiny-Inpaint",
+  "initial_inpaint_revision": "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
+  "teacher_model": "stable-diffusion-v1-5/stable-diffusion-inpainting",
+  "teacher_revision": "8a4288a76071f7280aedbdb3253bdb9e9d5d84bb",
+  "dataset": "prithivMLmods/Caption3o-Opt",
+  "dataset_revision": "17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b",
+  "dataset_split": "train",
+  "validation_samples": 128,
+  "image_column": "image",
+  "caption_column": "caption",
+  "resolution": 512,
   "unet_in_channels": 9,
   "mask_semantics": "white=regenerate, black=preserve",
-  "base_model": "neonforestmist/Clover-Image-Tiny",
-  "base_revision": "63b0e9f6be9c00888ff464f342a9ef052bf76681"
+  "mask_distribution": [
+    "brush",
+    "multi_brush",
+    "rectangle",
+    "ellipse",
+    "polygon",
+    "multi_region",
+    "outpaint"
+  ],
+  "training": {
+    "steps": 500,
+    "learning_rate": 0.000005,
+    "warmup_steps": 25,
+    "train_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "effective_batch_size": 4,
+    "mixed_precision": "bf16",
+    "gradient_checkpointing": true,
+    "caption_dropout_probability": 0.1,
+    "mask_min_area": 0.04,
+    "mask_max_area": 0.65,
+    "snr_gamma": 5.0,
+    "teacher_loss_weight": 0.75,
+    "ground_truth_loss_weight": 0.25,
+    "context_loss_weight": 0.25,
+    "masked_loss_weight": 2.5,
+    "boundary_loss_weight": 2.0,
+    "checkpointing_steps": 500,
+    "checkpoints_total_limit": 3,
+    "seed": 20260811
+  },
+  "modal": {
+    "profile": "guccichungus69",
+    "gpu": "A10",
+    "timeout_hours": 2,
+    "output_name": "clover-image-tiny-inpaint-v2-pilot-500"
+  },
+  "recommended_inference": {
+    "scheduler": "DPMSolverMultistepScheduler",
+    "steps": 20,
+    "guidance_scale": 6.0,
+    "mask_crop_padding": 96,
+    "retry_full_frame": true,
+    "composite_outside_mask": "exact source pixels"
+  }
 }

inpainting/README.md

@@ -1,6 +1,6 @@
 ---
 library_name: diffusers
-pipeline_tag: image-inpainting
+pipeline_tag: image-to-image
 base_model: neonforestmist/Clover-Image-Tiny
 license: creativeml-openrail-m
 tags:
@@ -8,14 +8,14 @@ tags:
   - inpainting
   - stable-diffusion
   - coreml
-  - iphone
 ---
 
 # Clover Image Tiny Inpaint 🍀
 
 An inpainting adaptation of Clover Image Tiny for 512×512 local generation and
 on-device Core ML deployment. White mask pixels are regenerated; black pixels
-are preserved.
+are preserved. The intended target is an SD 1.4-class architecture, rather than
+a device-specific model.
 
 The model uses a 9-channel U-Net input:
 
@@ -25,32 +25,54 @@ The model uses a 9-channel U-Net input:
 
 The base text encoder, VAE, scheduler, safety checker, and tokenizer remain
 compatible with Clover Image Tiny. The inpainting export additionally includes
-the VAE encoder needed to prepare the masked-image latent on iPhone.
+the VAE encoder needed to prepare the masked-image latent on device.
 
-## Diffusers
+## Diffusers example
 
 ```python
-from diffusers import AutoPipelineForInpainting
+from diffusers import AutoPipelineForInpainting, DPMSolverMultistepScheduler
 from diffusers.utils import load_image
 
 pipe = AutoPipelineForInpainting.from_pretrained(
     "neonforestmist/Clover-Image-Tiny-Inpaint",
     torch_dtype="auto",
 )
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 image = pipe(
     prompt="a tiny glass greenhouse glowing in a moonlit garden",
     image=load_image("input.png"),
     mask_image=load_image("mask.png"),
-    num_inference_steps=30,
+    num_inference_steps=20,
+    guidance_scale=6.0,
+    padding_mask_crop=96,
 ).images[0]
 image.save("clover-inpaint.png")
 ```
 
-## Core ML and iPhone 15
+The mask is a grayscale image: white means “regenerate” and black means
+“preserve.” A few useful inpainting prompts are:
+
+- `replace the masked area with a tiny glass greenhouse glowing at night`
+- `remove the person from the masked area and continue the background naturally`
+- `add a red enamel kettle on the masked countertop`
+
+Use DPM-Solver++, 20 steps, CFG 6.0, and a 96-pixel context margin as the
+interactive defaults. Exact-mask compositing preserves every source pixel
+outside the edit. The selected v2 checkpoint reduced deterministic held-out
+masked MAE from 0.3120 to 0.2517 and improved mean semantic CLIP alignment from
+0.2726 to 0.2890 over the previous release.
+
+For repeatable experiments, keep the input image, mask, seed, scheduler, and
+step count together. Inpainting is local editing: the unmasked region is
+provided as the masked-image conditioning and is also preserved by the native
+runtime compositor.
+
+## Core ML and SD 1.4-class deployment
 
 The companion Core ML resource bundle is converted for iOS 18 with a
-batch-one U-Net and chunked U-Net resources. The Swift runtime performs
-classifier-free guidance as two serial passes to reduce peak memory. The
+batch-one U-Net and chunked U-Net resources for the SD 1.4-class 512×512
+architecture. The Swift runtime performs classifier-free guidance as two
+serial passes to reduce peak memory. The
 bundled `VAEEncoder.mlmodelc` creates the masked-image latent locally, so the
 input image and mask do not leave the device.
 
@@ -62,7 +84,29 @@ Conversion and the native iOS integration live in the source Clover repo:
 
 ## Training provenance
 
-Training uses synthetic rectangle, ellipse, and brush masks over the pinned
-Apache-2.0 `prithivMLmods/Caption3o-Opt` image-caption dataset. The full job is
-launched by Modal under the `guccichungus69` workspace and stores its output in
-the `clover-image-tiny-inpaint-output` Volume before Core ML conversion.
+The context-aware v2 recipe distills the pinned official SD 1.5 inpainting
+U-Net into Clover's compact nine-channel U-Net while retaining a ground-truth
+diffusion loss. Training uses diverse synthetic free-form, multi-region,
+object-like, and outpainting masks over the pinned Apache-2.0
+`prithivMLmods/Caption3o-Opt` image-caption dataset. A deterministic 128-image
+holdout is excluded from optimization and reused for same-mask, same-seed
+comparisons against the previous release.
+
+The job runs on Modal under the `guccichungus69` workspace and stores rolling
+resumable checkpoints plus the final Diffusers pipeline in the
+`clover-image-tiny-inpaint-output` Volume before Core ML conversion. Exact
+revisions and objective weights are recorded in `inpainting/config.json` and
+the generated `training-summary.json`.
+
+## Citation
+
+If Clover Image Tiny Inpaint is useful in your work, please cite the release:
+
+```bibtex
+@software{lozadaperez2026cloverimagetinyinpaint,
+  author = {Lukas Lozada Perez},
+  title = {Clover Image Tiny Inpaint: Compact SD 1.4-Class Image Inpainting},
+  year = {2026},
+  url = {https://huggingface.co/neonforestmist/Clover-Image-Tiny-Inpaint}
+}
+```

inpainting/__init__.py

@@ -1,3 +1,3 @@
 """Clover Image Tiny inpainting training and export helpers."""
 
-__all__ = ["model", "masks"]
+__all__ = ["masks", "model"]

inpainting/config.json

@@ -1,29 +1,63 @@
 {
+  "recipe_version": 2,
   "base_model": "neonforestmist/Clover-Image-Tiny",
   "base_revision": "63b0e9f6be9c00888ff464f342a9ef052bf76681",
+  "initial_inpaint_model": "neonforestmist/Clover-Image-Tiny-Inpaint",
+  "initial_inpaint_revision": "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
+  "teacher_model": "stable-diffusion-v1-5/stable-diffusion-inpainting",
+  "teacher_revision": "8a4288a76071f7280aedbdb3253bdb9e9d5d84bb",
   "dataset": "prithivMLmods/Caption3o-Opt",
   "dataset_revision": "17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b",
   "dataset_split": "train",
+  "validation_samples": 128,
   "image_column": "image",
   "caption_column": "caption",
   "resolution": 512,
   "unet_in_channels": 9,
   "mask_semantics": "white=regenerate, black=preserve",
+  "mask_distribution": [
+    "brush",
+    "multi_brush",
+    "rectangle",
+    "ellipse",
+    "polygon",
+    "multi_region",
+    "outpaint"
+  ],
   "training": {
-    "steps": 4000,
-    "learning_rate": 0.00001,
-    "warmup_steps": 200,
+    "steps": 500,
+    "learning_rate": 0.000005,
+    "warmup_steps": 25,
     "train_batch_size": 1,
     "gradient_accumulation_steps": 4,
-    "mixed_precision": "fp16",
+    "effective_batch_size": 4,
+    "mixed_precision": "bf16",
     "gradient_checkpointing": true,
-    "mask_min_area": 0.12,
-    "mask_max_area": 0.55
+    "caption_dropout_probability": 0.1,
+    "mask_min_area": 0.04,
+    "mask_max_area": 0.65,
+    "snr_gamma": 5.0,
+    "teacher_loss_weight": 0.75,
+    "ground_truth_loss_weight": 0.25,
+    "context_loss_weight": 0.25,
+    "masked_loss_weight": 2.5,
+    "boundary_loss_weight": 2.0,
+    "checkpointing_steps": 500,
+    "checkpoints_total_limit": 3,
+    "seed": 20260811
   },
   "modal": {
     "profile": "guccichungus69",
-    "gpu": "A10G",
-    "timeout_hours": 4,
-    "estimated_hourly_usd": 1.0
+    "gpu": "A10",
+    "timeout_hours": 2,
+    "output_name": "clover-image-tiny-inpaint-v2-pilot-500"
+  },
+  "recommended_inference": {
+    "scheduler": "DPMSolverMultistepScheduler",
+    "steps": 20,
+    "guidance_scale": 6.0,
+    "mask_crop_padding": 96,
+    "retry_full_frame": true,
+    "composite_outside_mask": "exact source pixels"
   }
 }

inpainting/evaluate.py

@@ -0,0 +1,247 @@
+#!/usr/bin/env python3
+"""Compare two Clover inpainting checkpoints on held-out reconstruction masks."""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import random
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from diffusers import DPMSolverMultistepScheduler, StableDiffusionInpaintPipeline
+from PIL import Image, ImageDraw, ImageFilter, ImageOps
+
+from inpainting.masks import mask_area_fraction, random_mask
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--baseline_model", required=True)
+    parser.add_argument("--baseline_revision")
+    parser.add_argument("--candidate_model", required=True)
+    parser.add_argument("--dataset_name", required=True)
+    parser.add_argument("--dataset_revision")
+    parser.add_argument("--dataset_split", default="train")
+    parser.add_argument("--image_column", default="image")
+    parser.add_argument("--caption_column", default="caption")
+    parser.add_argument("--validation_samples", type=int, default=128)
+    parser.add_argument("--sample_count", type=int, default=6)
+    parser.add_argument("--steps", type=int, default=30)
+    parser.add_argument("--guidance_scale", type=float, default=7.5)
+    parser.add_argument("--seed", type=int, default=20260811)
+    parser.add_argument("--output_dir", type=Path, required=True)
+    return parser.parse_args()
+
+
+def _caption(value: Any) -> str:
+    if isinstance(value, list):
+        value = value[0] if value else ""
+    return " ".join(str(value or "").split())
+
+
+def _prepare_image(image: Image.Image, resolution: int = 512) -> Image.Image:
+    image = ImageOps.exif_transpose(image).convert("RGB")
+    side = min(image.size)
+    left = (image.width - side) // 2
+    top = (image.height - side) // 2
+    return image.crop((left, top, left + side, top + side)).resize(
+        (resolution, resolution), Image.Resampling.LANCZOS
+    )
+
+
+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.convert("RGB"), _feather_inside(mask))
+
+
+def _masked_mae(result: Image.Image, target: Image.Image, mask: Image.Image) -> float:
+    result_array = np.asarray(result, dtype=np.float32) / 255.0
+    target_array = np.asarray(target, dtype=np.float32) / 255.0
+    mask_array = np.asarray(mask.convert("L"), dtype=np.float32) / 255.0
+    denominator = max(1.0, mask_array.sum() * 3.0)
+    return float((np.abs(result_array - target_array) * mask_array[..., None]).sum() / denominator)
+
+
+def _black_fraction(image: Image.Image, mask: Image.Image) -> float:
+    pixels = np.asarray(image.convert("RGB"), dtype=np.uint8)
+    selected = np.asarray(mask.convert("L")) >= 128
+    if not selected.any():
+        return 1.0
+    return float(np.all(pixels[selected] <= 8, axis=1).mean())
+
+
+def _load_pipeline(model: str, *, revision: str | None = None):
+    kwargs = {"revision": revision} if revision else {}
+    pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+        model,
+        torch_dtype=torch.float16,
+        safety_checker=None,
+        requires_safety_checker=False,
+        **kwargs,
+    )
+    pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+        pipeline.scheduler.config,
+        algorithm_type="dpmsolver++",
+    )
+    return pipeline.to("cuda")
+
+
+def _run_model(
+    *,
+    model_name: str,
+    revision: str | None,
+    cases: list[dict[str, Any]],
+    steps: int,
+    guidance_scale: float,
+    output_dir: Path,
+) -> list[dict[str, Any]]:
+    pipeline = _load_pipeline(model_name, revision=revision)
+    records = []
+    for index, case in enumerate(cases):
+        generator = torch.Generator(device="cuda").manual_seed(case["seed"])
+        response = pipeline(
+            prompt=case["caption"],
+            image=case["source"],
+            mask_image=case["mask"],
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+            height=512,
+            width=512,
+        )
+        raw = response.images[0].convert("RGB")
+        result = _composite(raw, case["source"], case["mask"])
+        result.save(output_dir / f"{index:02d}.png")
+        records.append(
+            {
+                "index": index,
+                "masked_mae": _masked_mae(result, case["source"], case["mask"]),
+                "black_fraction": _black_fraction(result, case["mask"]),
+                "nsfw_content_detected": None,
+            }
+        )
+    del pipeline
+    gc.collect()
+    torch.cuda.empty_cache()
+    return records
+
+
+def _make_sheet(
+    cases: list[dict[str, Any]],
+    baseline_dir: Path,
+    candidate_dir: Path,
+    destination: Path,
+) -> None:
+    cell = 512
+    label_height = 32
+    sheet = Image.new("RGB", (cell * 3, (cell + label_height) * len(cases)), "white")
+    draw = ImageDraw.Draw(sheet)
+    for index, case in enumerate(cases):
+        row_y = index * (cell + label_height)
+        source = case["source"].copy()
+        overlay = Image.new("RGB", source.size, (255, 255, 255))
+        source = Image.blend(source, Image.composite(overlay, source, case["mask"]), 0.55)
+        images = [
+            source,
+            Image.open(baseline_dir / f"{index:02d}.png").convert("RGB"),
+            Image.open(candidate_dir / f"{index:02d}.png").convert("RGB"),
+        ]
+        for column, image in enumerate(images):
+            sheet.paste(image, (column * cell, row_y + label_height))
+        draw.text((8, row_y + 8), f"source + mask | {case['caption'][:58]}", fill="black")
+        draw.text((cell + 8, row_y + 8), "current checkpoint", fill="black")
+        draw.text((cell * 2 + 8, row_y + 8), "candidate", fill="black")
+    sheet.save(destination)
+
+
+def main() -> None:
+    args = parse_args()
+    args.output_dir.mkdir(parents=True, exist_ok=False)
+    baseline_dir = args.output_dir / "baseline"
+    candidate_dir = args.output_dir / "candidate"
+    baseline_dir.mkdir()
+    candidate_dir.mkdir()
+
+    dataset = load_dataset(
+        args.dataset_name,
+        split=args.dataset_split,
+        revision=args.dataset_revision,
+    ).shuffle(seed=args.seed)
+    if args.validation_samples < args.sample_count or args.validation_samples >= len(dataset):
+        raise ValueError("validation set must contain at least sample_count records")
+    validation = dataset.select(
+        range(len(dataset) - args.validation_samples, len(dataset))
+    ).select(range(args.sample_count))
+
+    cases = []
+    for index, example in enumerate(validation):
+        image = example[args.image_column]
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(np.asarray(image))
+        source = _prepare_image(image)
+        rng = random.Random(args.seed + index * 1009)
+        mask = random_mask((512, 512), rng, min_area=0.08, max_area=0.45)
+        caption = _caption(example[args.caption_column])
+        cases.append(
+            {
+                "source": source,
+                "mask": mask,
+                "caption": caption,
+                "seed": args.seed + index,
+                "mask_area": mask_area_fraction(mask),
+            }
+        )
+        source.save(args.output_dir / f"source-{index:02d}.png")
+        mask.save(args.output_dir / f"mask-{index:02d}.png")
+
+    baseline_records = _run_model(
+        model_name=args.baseline_model,
+        revision=args.baseline_revision,
+        cases=cases,
+        steps=args.steps,
+        guidance_scale=args.guidance_scale,
+        output_dir=baseline_dir,
+    )
+    candidate_records = _run_model(
+        model_name=args.candidate_model,
+        revision=None,
+        cases=cases,
+        steps=args.steps,
+        guidance_scale=args.guidance_scale,
+        output_dir=candidate_dir,
+    )
+    _make_sheet(cases, baseline_dir, candidate_dir, args.output_dir / "comparison.png")
+
+    metrics = {
+        "baseline_model": args.baseline_model,
+        "baseline_revision": args.baseline_revision,
+        "candidate_model": args.candidate_model,
+        "dataset": args.dataset_name,
+        "dataset_revision": args.dataset_revision,
+        "sample_count": args.sample_count,
+        "steps": args.steps,
+        "baseline": baseline_records,
+        "candidate": candidate_records,
+        "baseline_mean_masked_mae": float(
+            np.mean([record["masked_mae"] for record in baseline_records])
+        ),
+        "candidate_mean_masked_mae": float(
+            np.mean([record["masked_mae"] for record in candidate_records])
+        ),
+    }
+    (args.output_dir / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+    print(json.dumps(metrics, indent=2))
+
+
+if __name__ == "__main__":
+    main()

inpainting/evaluate_semantic.py

@@ -0,0 +1,494 @@
+#!/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()

inpainting/masks.py

@@ -1,57 +1,195 @@
-"""Synthetic mask generation used for inpainting fine-tuning."""
+"""Diverse synthetic masks used for context-aware inpainting training."""
 
 from __future__ import annotations
 
+import math
 import random
 
-from PIL import Image, ImageChops, ImageDraw
+from PIL import Image, ImageChops, ImageDraw, ImageFilter
 
+MASK_KINDS = (
+    "brush",
+    "multi_brush",
+    "rectangle",
+    "ellipse",
+    "polygon",
+    "multi_region",
+    "outpaint",
+)
 
-def random_mask(
+
+def _area_fraction(mask: Image.Image) -> float:
+    histogram = mask.convert("L").histogram()
+    white_sum = sum(value * count for value, count in enumerate(histogram))
+    return white_sum / (255.0 * mask.width * mask.height)
+
+
+def _random_box(
+    size: tuple[int, int],
+    rng: random.Random,
+    target_area: float,
+) -> tuple[int, int, int, int]:
+    width, height = size
+    aspect = math.exp(rng.uniform(math.log(0.35), math.log(2.85)))
+    box_width = min(width, max(8, round(math.sqrt(target_area * aspect))))
+    box_height = min(height, max(8, round(math.sqrt(target_area / aspect))))
+    left = rng.randint(0, max(0, width - box_width))
+    top = rng.randint(0, max(0, height - box_height))
+    return left, top, left + box_width, top + box_height
+
+
+def _draw_brush(
+    draw: ImageDraw.ImageDraw,
     size: tuple[int, int],
     rng: random.Random,
     *,
-    min_area: float = 0.12,
-    max_area: float = 0.55,
-) -> Image.Image:
-    """Return a black/white L mask where white means "paint this area"."""
+    strokes: int,
+) -> None:
+    width, height = size
+    for _ in range(strokes):
+        stroke_width = max(6, round(min(width, height) * rng.uniform(0.035, 0.16)))
+        point_count = rng.randint(3, 8)
+        x = rng.uniform(0, width - 1)
+        y = rng.uniform(0, height - 1)
+        points: list[tuple[float, float]] = [(x, y)]
+        angle = rng.uniform(0, 2 * math.pi)
+        for _ in range(point_count - 1):
+            angle += rng.uniform(-1.25, 1.25)
+            distance = rng.uniform(0.06, 0.24) * min(width, height)
+            x = min(width - 1, max(0, x + math.cos(angle) * distance))
+            y = min(height - 1, max(0, y + math.sin(angle) * distance))
+            points.append((x, y))
+        draw.line(points, fill=255, width=stroke_width, joint="curve")
+        radius = stroke_width / 2
+        for px, py in points:
+            draw.ellipse((px - radius, py - radius, px + radius, py + radius), fill=255)
+
+
+def _draw_outpaint(
+    draw: ImageDraw.ImageDraw,
+    size: tuple[int, int],
+    rng: random.Random,
+    target_fraction: float,
+) -> None:
+    width, height = size
+    sides = rng.sample(("left", "right", "top", "bottom"), rng.choice((1, 1, 1, 2)))
+    remaining = max(0.02, target_fraction)
+    for index, side in enumerate(sides):
+        share = remaining if index == len(sides) - 1 else remaining * rng.uniform(0.35, 0.7)
+        if side in {"left", "right"}:
+            thickness = max(4, min(width, round(share * width)))
+            box = (0, 0, thickness, height) if side == "left" else (width - thickness, 0, width, height)
+        else:
+            thickness = max(4, min(height, round(share * height)))
+            box = (0, 0, width, thickness) if side == "top" else (0, height - thickness, width, height)
+        draw.rectangle(box, fill=255)
+        remaining = max(0.0, remaining - share)
+
 
+def _draw_candidate(
+    size: tuple[int, int],
+    rng: random.Random,
+    *,
+    kind: str,
+    target_fraction: float,
+) -> Image.Image:
     width, height = size
     mask = Image.new("L", size, 0)
     draw = ImageDraw.Draw(mask)
-    shape = rng.choice(("rectangle", "ellipse", "brush"))
-
-    target_area = rng.uniform(min_area, max_area) * width * height
-    if shape in {"rectangle", "ellipse"}:
-        aspect = rng.uniform(0.45, 2.2)
-        box_width = max(8, int((target_area * aspect) ** 0.5))
-        box_height = max(8, int((target_area / aspect) ** 0.5))
-        box_width = min(box_width, width)
-        box_height = min(box_height, height)
-        left = rng.randint(0, max(0, width - box_width))
-        top = rng.randint(0, max(0, height - box_height))
-        box = (left, top, left + box_width, top + box_height)
-        if shape == "rectangle":
-            draw.rectangle(box, fill=255)
-        else:
-            draw.ellipse(box, fill=255)
-    else:
-        # A few overlapping strokes cover irregular object-shaped regions while
-        # staying cheap enough to generate inside a DataLoader worker.
-        stroke_width = max(8, int(min(width, height) * rng.uniform(0.06, 0.18)))
-        points = [
-            (
-                rng.randint(0, width - 1),
-                rng.randint(0, height - 1),
+    target_area = target_fraction * width * height
+
+    if kind == "rectangle":
+        draw.rounded_rectangle(
+            _random_box(size, rng, target_area),
+            radius=rng.randint(0, max(1, round(min(width, height) * 0.08))),
+            fill=255,
+        )
+    elif kind == "ellipse":
+        draw.ellipse(_random_box(size, rng, target_area), fill=255)
+    elif kind == "polygon":
+        center_x = rng.uniform(width * 0.2, width * 0.8)
+        center_y = rng.uniform(height * 0.2, height * 0.8)
+        radius = math.sqrt(target_area / math.pi)
+        points = []
+        point_count = rng.randint(5, 10)
+        for index in range(point_count):
+            angle = (2 * math.pi * index / point_count) + rng.uniform(-0.25, 0.25)
+            local_radius = radius * rng.uniform(0.65, 1.35)
+            points.append(
+                (
+                    min(width - 1, max(0, center_x + math.cos(angle) * local_radius)),
+                    min(height - 1, max(0, center_y + math.sin(angle) * local_radius)),
+                )
             )
-            for _ in range(rng.randint(2, 5))
-        ]
-        draw.line(points, fill=255, width=stroke_width, joint="curve")
-        radius = stroke_width // 2
-        for x, y in points:
-            draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=255)
+        draw.polygon(points, fill=255)
+    elif kind == "brush":
+        _draw_brush(draw, size, rng, strokes=1)
+    elif kind == "multi_brush":
+        _draw_brush(draw, size, rng, strokes=rng.randint(2, 5))
+    elif kind == "multi_region":
+        region_count = rng.randint(2, 5)
+        for _ in range(region_count):
+            region_area = target_area * rng.uniform(0.12, 0.5)
+            box = _random_box(size, rng, region_area)
+            if rng.random() < 0.55:
+                draw.ellipse(box, fill=255)
+            else:
+                draw.rounded_rectangle(box, radius=rng.randint(2, 24), fill=255)
+    elif kind == "outpaint":
+        _draw_outpaint(draw, size, rng, target_fraction)
+    else:
+        raise ValueError(f"Unsupported mask kind: {kind}")
 
-    return mask
+    if kind not in {"rectangle", "outpaint"} and rng.random() < 0.35:
+        # A small close operation removes pinholes without making every edge
+        # unnaturally geometric.
+        mask = mask.filter(ImageFilter.MaxFilter(rng.choice((3, 5, 7))))
+    return mask.point(lambda value: 255 if value >= 128 else 0, mode="L")
+
+
+def random_mask(
+    size: tuple[int, int],
+    rng: random.Random,
+    *,
+    min_area: float = 0.04,
+    max_area: float = 0.65,
+    kind: str | None = None,
+) -> Image.Image:
+    """Return a binary mask where white means regenerate.
+
+    The sampler deliberately mixes object-like regions, free-form user brush
+    strokes, disconnected edits, and edge/outpainting masks. Candidates are
+    retried so the actual white area—not only a geometric estimate—falls close
+    to the requested range.
+    """
+
+    if not 0.0 < min_area < max_area < 1.0:
+        raise ValueError("mask area bounds must satisfy 0 < min < max < 1")
+    if kind is not None and kind not in MASK_KINDS:
+        raise ValueError(f"Unsupported mask kind: {kind}")
+
+    best_mask: Image.Image | None = None
+    best_distance = float("inf")
+    target = rng.uniform(min_area, max_area)
+    for _ in range(16):
+        selected_kind = kind or rng.choices(
+            MASK_KINDS,
+            weights=(24, 18, 12, 10, 12, 16, 8),
+            k=1,
+        )[0]
+        candidate = _draw_candidate(size, rng, kind=selected_kind, target_fraction=target)
+        area = _area_fraction(candidate)
+        if min_area <= area <= max_area:
+            return candidate
+        distance = min(abs(area - min_area), abs(area - max_area))
+        if distance < best_distance:
+            best_mask = candidate
+            best_distance = distance
+        target = rng.uniform(min_area, max_area)
+
+    if best_mask is None:
+        raise RuntimeError("mask generation did not produce a candidate")
+    return best_mask
 
 
 def apply_mask(image: Image.Image, mask: Image.Image) -> Image.Image:
@@ -60,3 +198,9 @@ def apply_mask(image: Image.Image, mask: Image.Image) -> Image.Image:
     image = image.convert("RGB")
     keep = ImageChops.invert(mask.convert("L"))
     return Image.composite(image, Image.new("RGB", image.size), keep)
+
+
+def mask_area_fraction(mask: Image.Image) -> float:
+    """Expose the exact white-area fraction for validation and reporting."""
+
+    return _area_fraction(mask)

inpainting/objective.py

@@ -0,0 +1,71 @@
+"""Loss helpers for context-aware inpainting distillation."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+def min_snr_weights(
+    alphas_cumprod: torch.Tensor,
+    timesteps: torch.Tensor,
+    *,
+    gamma: float,
+    prediction_type: str,
+) -> torch.Tensor:
+    """Return the Min-SNR weighting from diffusion fine-tuning literature."""
+
+    alpha = alphas_cumprod.to(device=timesteps.device, dtype=torch.float32)[timesteps]
+    snr = alpha / (1.0 - alpha).clamp_min(1e-8)
+    clipped = torch.minimum(snr, torch.full_like(snr, gamma))
+    if prediction_type == "epsilon":
+        return clipped / snr.clamp_min(1e-8)
+    if prediction_type == "v_prediction":
+        return clipped / (snr + 1.0)
+    raise ValueError(f"Unsupported prediction type: {prediction_type}")
+
+
+def spatial_loss_weights(
+    mask: torch.Tensor,
+    *,
+    context_weight: float,
+    masked_weight: float,
+    boundary_weight: float,
+    boundary_radius: int = 2,
+) -> torch.Tensor:
+    """Emphasize regenerated pixels and the context-sensitive mask boundary."""
+
+    if mask.ndim != 4 or mask.shape[1] != 1:
+        raise ValueError(f"Expected BCHW single-channel mask, got {tuple(mask.shape)}")
+    if boundary_radius < 1:
+        raise ValueError("boundary_radius must be positive")
+    binary = (mask >= 0.5).to(dtype=torch.float32)
+    kernel = boundary_radius * 2 + 1
+    dilated = F.max_pool2d(binary, kernel_size=kernel, stride=1, padding=boundary_radius)
+    eroded = 1.0 - F.max_pool2d(
+        1.0 - binary,
+        kernel_size=kernel,
+        stride=1,
+        padding=boundary_radius,
+    )
+    boundary = (dilated - eroded).clamp(0.0, 1.0)
+    weights = torch.full_like(binary, float(context_weight))
+    weights = weights + binary * (float(masked_weight) - float(context_weight))
+    weights = weights + boundary * float(boundary_weight)
+    return weights / weights.mean(dim=(1, 2, 3), keepdim=True).clamp_min(1e-8)
+
+
+def weighted_mse(
+    prediction: torch.Tensor,
+    target: torch.Tensor,
+    *,
+    spatial_weights: torch.Tensor,
+    sample_weights: torch.Tensor,
+) -> torch.Tensor:
+    """Compute channel-averaged, spatially and per-sample weighted MSE."""
+
+    if prediction.shape != target.shape:
+        raise ValueError("prediction and target must have identical shapes")
+    per_pixel = (prediction.float() - target.float()).square().mean(dim=1, keepdim=True)
+    per_sample = (per_pixel * spatial_weights.float()).mean(dim=(1, 2, 3))
+    return (per_sample * sample_weights.float()).mean()

inpainting/test_v2.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import random
+import unittest
+
+import numpy as np
+import torch
+from PIL import Image
+
+from inpainting.masks import MASK_KINDS, apply_mask, mask_area_fraction, random_mask
+from inpainting.objective import min_snr_weights, spatial_loss_weights, weighted_mse
+
+
+class MaskTests(unittest.TestCase):
+    def test_every_mask_family_is_binary_and_nonempty(self) -> None:
+        for index, kind in enumerate(MASK_KINDS):
+            mask = random_mask(
+                (256, 256),
+                random.Random(1000 + index),
+                min_area=0.03,
+                max_area=0.70,
+                kind=kind,
+            )
+            values = set(np.unique(np.asarray(mask)).tolist())
+            self.assertTrue(values <= {0, 255}, (kind, values))
+            area = mask_area_fraction(mask)
+            self.assertGreater(area, 0.005, kind)
+            self.assertLess(area, 0.90, kind)
+
+    def test_apply_mask_only_blacks_white_region(self) -> None:
+        image = Image.new("RGB", (8, 8), (20, 40, 60))
+        mask = Image.new("L", (8, 8), 0)
+        for y in range(2, 6):
+            for x in range(3, 7):
+                mask.putpixel((x, y), 255)
+        result = np.asarray(apply_mask(image, mask))
+        self.assertTrue(np.all(result[2:6, 3:7] == 0))
+        self.assertTrue(np.all(result[0, 0] == (20, 40, 60)))
+
+
+class ObjectiveTests(unittest.TestCase):
+    def test_spatial_weights_prioritize_mask_and_boundary(self) -> None:
+        mask = torch.zeros(1, 1, 16, 16)
+        mask[:, :, 5:11, 5:11] = 1
+        weights = spatial_loss_weights(
+            mask,
+            context_weight=0.25,
+            masked_weight=2.5,
+            boundary_weight=2.0,
+            boundary_radius=1,
+        )
+        self.assertAlmostEqual(weights.mean().item(), 1.0, places=5)
+        self.assertGreater(weights[0, 0, 5, 5], weights[0, 0, 8, 8])
+        self.assertGreater(weights[0, 0, 8, 8], weights[0, 0, 0, 0])
+
+    def test_min_snr_and_weighted_mse_are_finite(self) -> None:
+        alphas = torch.linspace(0.999, 0.001, 1000)
+        timesteps = torch.tensor([0, 250, 999])
+        sample_weights = min_snr_weights(
+            alphas,
+            timesteps,
+            gamma=5.0,
+            prediction_type="epsilon",
+        )
+        prediction = torch.ones(3, 4, 8, 8)
+        target = torch.zeros_like(prediction)
+        spatial = torch.ones(3, 1, 8, 8)
+        loss = weighted_mse(
+            prediction,
+            target,
+            spatial_weights=spatial,
+            sample_weights=sample_weights,
+        )
+        self.assertTrue(torch.isfinite(loss))
+        self.assertGreater(loss.item(), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

inpainting/train.py

@@ -1,8 +1,10 @@
 #!/usr/bin/env python3
-"""Fine-tune Clover Image Tiny into a 9-channel inpainting pipeline.
+"""Train Clover Image Tiny v2 as a context-aware 9-channel inpainting model.
 
-The dataset only needs an image column and a text-caption column. Masks are
-generated on the fly, so one image produces many distinct training examples.
+The trainer warm-starts from the existing Clover inpainting checkpoint and
+distills the official SD 1.5 inpainting U-Net while retaining a ground-truth
+diffusion objective. Masks are synthesized on the fly with free-form,
+multi-region, object-like, and outpainting geometry.
 """
 
 from __future__ import annotations
@@ -14,6 +16,7 @@ import shutil
 from pathlib import Path
 from typing import Any
 
+import numpy as np
 import torch
 import torch.nn.functional as F
 from accelerate import Accelerator
@@ -22,20 +25,25 @@ from diffusers import (
     AutoencoderKL,
     DDPMScheduler,
     StableDiffusionInpaintPipeline,
+    UNet2DConditionModel,
     get_scheduler,
 )
-from PIL import Image
+from PIL import Image, ImageOps
 from transformers import CLIPTextModel, CLIPTokenizer
-from torchvision import transforms
 
-from inpainting.masks import apply_mask, random_mask
+from inpainting.masks import apply_mask, mask_area_fraction, random_mask
 from inpainting.model import make_inpainting_unet
+from inpainting.objective import min_snr_weights, spatial_loss_weights, weighted_mse
 
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--pretrained_model_name_or_path", required=True)
     parser.add_argument("--revision")
+    parser.add_argument("--initial_inpaint_model")
+    parser.add_argument("--initial_inpaint_revision")
+    parser.add_argument("--teacher_model_name_or_path")
+    parser.add_argument("--teacher_revision")
     parser.add_argument("--dataset_name", required=True)
     parser.add_argument("--dataset_config_name")
     parser.add_argument("--dataset_revision")
@@ -44,19 +52,32 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--caption_column", default="caption")
     parser.add_argument("--resolution", type=int, default=512)
     parser.add_argument("--train_batch_size", type=int, default=1)
-    parser.add_argument("--max_train_steps", type=int, default=4000)
-    parser.add_argument("--learning_rate", type=float, default=1e-5)
+    parser.add_argument("--max_train_steps", type=int, default=12000)
+    parser.add_argument("--learning_rate", type=float, default=5e-6)
     parser.add_argument("--lr_scheduler", default="cosine")
-    parser.add_argument("--lr_warmup_steps", type=int, default=200)
+    parser.add_argument("--lr_warmup_steps", type=int, default=500)
     parser.add_argument("--gradient_accumulation_steps", type=int, default=4)
     parser.add_argument("--gradient_checkpointing", action="store_true")
-    parser.add_argument("--mixed_precision", choices=("no", "fp16", "bf16"), default="fp16")
-    parser.add_argument("--seed", type=int, default=20260810)
-    parser.add_argument("--mask_min_area", type=float, default=0.12)
-    parser.add_argument("--mask_max_area", type=float, default=0.55)
+    parser.add_argument("--mixed_precision", choices=("no", "fp16", "bf16"), default="bf16")
+    parser.add_argument("--seed", type=int, default=20260811)
+    parser.add_argument("--mask_min_area", type=float, default=0.04)
+    parser.add_argument("--mask_max_area", type=float, default=0.65)
+    parser.add_argument("--caption_dropout_probability", type=float, default=0.10)
+    parser.add_argument("--teacher_loss_weight", type=float, default=0.75)
+    parser.add_argument("--ground_truth_loss_weight", type=float, default=0.25)
+    parser.add_argument("--context_loss_weight", type=float, default=0.25)
+    parser.add_argument("--masked_loss_weight", type=float, default=2.5)
+    parser.add_argument("--boundary_loss_weight", type=float, default=2.0)
+    parser.add_argument("--boundary_radius", type=int, default=2)
+    parser.add_argument("--snr_gamma", type=float, default=5.0)
+    parser.add_argument("--random_flip", action="store_true")
     parser.add_argument("--max_train_samples", type=int)
+    parser.add_argument("--validation_samples", type=int, default=128)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--checkpointing_steps", type=int, default=500)
+    parser.add_argument("--checkpoints_total_limit", type=int, default=3)
+    parser.add_argument("--resume_from_checkpoint")
     parser.add_argument("--output_dir", type=Path, required=True)
-    parser.add_argument("--validation_prompt")
     parser.add_argument("--push_to_hub", action="store_true")
     parser.add_argument("--hub_model_id")
     return parser.parse_args()
@@ -66,16 +87,43 @@ def model_kwargs(revision: str | None) -> dict[str, Any]:
     return {"revision": revision} if revision else {}
 
 
-def image_to_tensor(image: Image.Image, resolution: int) -> torch.Tensor:
-    transform = transforms.Compose(
-        [
-            transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR),
-            transforms.CenterCrop(resolution),
-            transforms.ToTensor(),
-            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
-        ]
-    )
-    return transform(image.convert("RGB"))
+def _frozen_weight_dtype(mixed_precision: str) -> torch.dtype:
+    if mixed_precision == "fp16":
+        return torch.float16
+    if mixed_precision == "bf16":
+        return torch.bfloat16
+    return torch.float32
+
+
+def _normalize_caption(value: Any) -> str:
+    if isinstance(value, list):
+        value = value[0] if value else ""
+    return " ".join(str(value or "").split())
+
+
+def _crop_image(
+    image: Image.Image,
+    *,
+    resolution: int,
+    rng: random.Random,
+    random_flip: bool,
+) -> Image.Image:
+    image = ImageOps.exif_transpose(image).convert("RGB")
+    width, height = image.size
+    scale = rng.uniform(0.78, 1.0)
+    crop_side = max(1, round(min(width, height) * scale))
+    left = rng.randint(0, max(0, width - crop_side))
+    top = rng.randint(0, max(0, height - crop_side))
+    image = image.crop((left, top, left + crop_side, top + crop_side))
+    image = image.resize((resolution, resolution), Image.Resampling.LANCZOS)
+    if random_flip and rng.random() < 0.5:
+        image = ImageOps.mirror(image)
+    return image
+
+
+def _image_to_tensor(image: Image.Image) -> torch.Tensor:
+    array = np.asarray(image.convert("RGB"), dtype=np.float32) / 127.5 - 1.0
+    return torch.from_numpy(array).permute(2, 0, 1).contiguous()
 
 
 def make_collate_fn(
@@ -85,25 +133,34 @@ def make_collate_fn(
     min_area: float,
     max_area: float,
     seed: int,
+    caption_dropout_probability: float,
+    random_flip: bool,
 ):
-    worker_rng = random.Random(seed)
+    worker_rng: random.Random | None = None
 
     def collate(examples: list[dict[str, Any]]) -> dict[str, Any]:
+        nonlocal worker_rng
+        if worker_rng is None:
+            # torch.initial_seed differs across DataLoader workers and remains
+            # stable for the lifetime of each worker process.
+            worker_rng = random.Random(seed ^ torch.initial_seed())
+
         images: list[torch.Tensor] = []
         masked_images: list[torch.Tensor] = []
         masks: list[torch.Tensor] = []
         captions: list[str] = []
+        mask_areas: list[float] = []
 
         for example in examples:
             image = example["image"]
             if not isinstance(image, Image.Image):
-                image = Image.fromarray(image)
-            image = image.convert("RGB")
-            image = transforms.Resize(
-                resolution,
-                interpolation=transforms.InterpolationMode.BILINEAR,
-            )(image)
-            image = transforms.CenterCrop(resolution)(image)
+                image = Image.fromarray(np.asarray(image))
+            image = _crop_image(
+                image,
+                resolution=resolution,
+                rng=worker_rng,
+                random_flip=random_flip,
+            )
             mask = random_mask(
                 (resolution, resolution),
                 worker_rng,
@@ -112,16 +169,16 @@ def make_collate_fn(
             )
             masked = apply_mask(image, mask)
 
-            images.append(image_to_tensor(image, resolution))
-            masked_images.append(image_to_tensor(masked, resolution))
-            mask_tensor = torch.from_numpy(
-                __import__("numpy").array(mask, dtype="float32") / 255.0
-            ).unsqueeze(0)
-            masks.append(mask_tensor)
-            caption = example["caption"]
-            if isinstance(caption, list):
-                caption = caption[0] if caption else ""
-            captions.append(str(caption or ""))
+            images.append(_image_to_tensor(image))
+            masked_images.append(_image_to_tensor(masked))
+            masks.append(
+                torch.from_numpy(np.asarray(mask, dtype=np.float32) / 255.0).unsqueeze(0)
+            )
+            caption = _normalize_caption(example["caption"])
+            if worker_rng.random() < caption_dropout_probability:
+                caption = ""
+            captions.append(caption)
+            mask_areas.append(mask_area_fraction(mask))
 
         tokenized = tokenizer(
             captions,
@@ -135,32 +192,56 @@ def make_collate_fn(
             "masked_pixel_values": torch.stack(masked_images),
             "mask": torch.stack(masks),
             "input_ids": tokenized.input_ids,
+            "mask_area": torch.tensor(mask_areas, dtype=torch.float32),
         }
 
     return collate
 
 
-def save_pipeline(
+def _load_student(args: argparse.Namespace) -> UNet2DConditionModel:
+    if args.initial_inpaint_model:
+        student = UNet2DConditionModel.from_pretrained(
+            args.initial_inpaint_model,
+            subfolder="unet",
+            low_cpu_mem_usage=False,
+            **model_kwargs(args.initial_inpaint_revision),
+        )
+        if student.config.in_channels != 9:
+            raise ValueError("initial inpainting U-Net must have nine input channels")
+        return student
+    return make_inpainting_unet(
+        args.pretrained_model_name_or_path,
+        revision=args.revision,
+    )
+
+
+def _save_pipeline(
     *,
-    base_model: str,
-    revision: str | None,
-    unet,
+    args: argparse.Namespace,
+    unet: UNet2DConditionModel,
     output_dir: Path,
 ) -> None:
     output_dir.mkdir(parents=True, exist_ok=True)
     pipeline = StableDiffusionInpaintPipeline.from_pretrained(
-        base_model,
-        revision=revision,
+        args.pretrained_model_name_or_path,
+        revision=args.revision,
         unet=unet,
     )
     pipeline.save_pretrained(output_dir, safe_serialization=True)
     (output_dir / "inpainting-config.json").write_text(
         json.dumps(
             {
+                "recipe_version": 2,
                 "unet_in_channels": 9,
                 "mask_semantics": "white=regenerate, black=preserve",
-                "base_model": base_model,
-                "base_revision": revision,
+                "base_model": args.pretrained_model_name_or_path,
+                "base_revision": args.revision,
+                "initial_inpaint_model": args.initial_inpaint_model,
+                "initial_inpaint_revision": args.initial_inpaint_revision,
+                "teacher_model": args.teacher_model_name_or_path,
+                "teacher_revision": args.teacher_revision,
+                "objective": "teacher distillation + ground-truth diffusion",
+                "mask_distribution": "brush, multi-brush, box, ellipse, polygon, multi-region, outpaint",
             },
             indent=2,
         )
@@ -168,28 +249,60 @@ def save_pipeline(
     )
 
 
-def main() -> None:
-    args = parse_args()
+def _prune_checkpoints(checkpoint_root: Path, limit: int) -> None:
+    checkpoints = sorted(
+        (path for path in checkpoint_root.glob("checkpoint-*") if path.is_dir()),
+        key=lambda path: int(path.name.rsplit("-", 1)[-1]),
+    )
+    for path in checkpoints[:-limit]:
+        shutil.rmtree(path)
+
+
+def _validate_args(args: argparse.Namespace) -> None:
     if not 0.0 < args.mask_min_area < args.mask_max_area < 1.0:
         raise ValueError("mask area bounds must satisfy 0 < min < max < 1")
+    if not 0.0 <= args.caption_dropout_probability < 1.0:
+        raise ValueError("caption dropout probability must be in [0, 1)")
+    if args.teacher_loss_weight < 0 or args.ground_truth_loss_weight < 0:
+        raise ValueError("loss weights must be non-negative")
+    if args.teacher_loss_weight > 0 and not args.teacher_model_name_or_path:
+        raise ValueError("teacher model is required when teacher loss is enabled")
+    if args.teacher_loss_weight + args.ground_truth_loss_weight <= 0:
+        raise ValueError("at least one training objective must be enabled")
     if args.push_to_hub and not args.hub_model_id:
         raise ValueError("--hub_model_id is required with --push_to_hub")
+    if args.resolution % 8:
+        raise ValueError("resolution must be divisible by 8")
+
+
+def main() -> None:
+    args = parse_args()
+    _validate_args(args)
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint_root = args.output_dir / "checkpoints"
+    metrics_path = args.output_dir / "training-metrics.jsonl"
 
     accelerator = Accelerator(
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision=args.mixed_precision,
     )
-    accelerator.init_trackers("clover-image-tiny-inpaint")
     torch.manual_seed(args.seed)
     random.seed(args.seed)
+    np.random.seed(args.seed % (2**32))
 
     kwargs = model_kwargs(args.revision)
+    weight_dtype = _frozen_weight_dtype(args.mixed_precision)
     dataset = load_dataset(
         args.dataset_name,
         args.dataset_config_name,
         split=args.dataset_split,
         revision=args.dataset_revision,
     )
+    dataset = dataset.shuffle(seed=args.seed)
+    if args.validation_samples < 0 or args.validation_samples >= len(dataset):
+        raise ValueError("validation_samples must be smaller than the dataset")
+    if args.validation_samples:
+        dataset = dataset.select(range(len(dataset) - args.validation_samples))
     if args.max_train_samples:
         dataset = dataset.select(range(min(args.max_train_samples, len(dataset))))
     if args.image_column not in dataset.column_names:
@@ -208,29 +321,41 @@ def main() -> None:
     text_encoder = CLIPTextModel.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="text_encoder",
+        dtype=weight_dtype,
         **kwargs,
     )
     vae = AutoencoderKL.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="vae",
+        torch_dtype=weight_dtype,
         **kwargs,
     )
-    unet = make_inpainting_unet(
-        args.pretrained_model_name_or_path,
-        revision=args.revision,
-    )
+    student = _load_student(args)
     noise_scheduler = DDPMScheduler.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="scheduler",
         **kwargs,
     )
+    teacher = None
+    if args.teacher_loss_weight > 0:
+        teacher = UNet2DConditionModel.from_pretrained(
+            args.teacher_model_name_or_path,
+            subfolder="unet",
+            low_cpu_mem_usage=True,
+            torch_dtype=weight_dtype,
+            use_safetensors=True,
+            variant="fp16",
+            **model_kwargs(args.teacher_revision),
+        )
+        if teacher.config.in_channels != 9:
+            raise ValueError("teacher inpainting U-Net must have nine input channels")
 
     if args.gradient_checkpointing:
-        unet.enable_gradient_checkpointing()
-    text_encoder.requires_grad_(False)
-    vae.requires_grad_(False)
-    text_encoder.eval()
-    vae.eval()
+        student.enable_gradient_checkpointing()
+    text_encoder.requires_grad_(False).eval()
+    vae.requires_grad_(False).eval()
+    if teacher is not None:
+        teacher.requires_grad_(False).eval()
 
     collate_fn = make_collate_fn(
         tokenizer=tokenizer,
@@ -238,17 +363,20 @@ def main() -> None:
         min_area=args.mask_min_area,
         max_area=args.mask_max_area,
         seed=args.seed,
+        caption_dropout_probability=args.caption_dropout_probability,
+        random_flip=args.random_flip,
     )
     dataloader = torch.utils.data.DataLoader(
         dataset,
         shuffle=True,
         collate_fn=collate_fn,
         batch_size=args.train_batch_size,
-        num_workers=2,
+        num_workers=args.num_workers,
         pin_memory=True,
+        persistent_workers=args.num_workers > 0,
     )
     optimizer = torch.optim.AdamW(
-        unet.parameters(),
+        student.parameters(),
         lr=args.learning_rate,
         betas=(0.9, 0.999),
         weight_decay=1e-2,
@@ -261,18 +389,42 @@ def main() -> None:
         num_training_steps=args.max_train_steps * accelerator.num_processes,
     )
 
-    unet, optimizer, dataloader, lr_scheduler = accelerator.prepare(
-        unet, optimizer, dataloader, lr_scheduler
+    student, optimizer, dataloader, lr_scheduler = accelerator.prepare(
+        student, optimizer, dataloader, lr_scheduler
     )
-    weight_dtype = torch.float32
-    if accelerator.mixed_precision == "fp16":
-        weight_dtype = torch.float16
-    elif accelerator.mixed_precision == "bf16":
-        weight_dtype = torch.bfloat16
-    vae.to(accelerator.device, dtype=weight_dtype)
-    text_encoder.to(accelerator.device, dtype=weight_dtype)
+    vae.to(accelerator.device)
+    text_encoder.to(accelerator.device)
+    if teacher is not None:
+        teacher.to(accelerator.device)
 
     global_step = 0
+    if args.resume_from_checkpoint:
+        checkpoint = Path(args.resume_from_checkpoint)
+        if args.resume_from_checkpoint == "latest":
+            candidates = sorted(
+                checkpoint_root.glob("checkpoint-*"),
+                key=lambda path: int(path.name.rsplit("-", 1)[-1]),
+            )
+            if not candidates:
+                raise ValueError("No checkpoint is available to resume")
+            checkpoint = candidates[-1]
+        accelerator.load_state(checkpoint)
+        global_step = int(checkpoint.name.rsplit("-", 1)[-1])
+        accelerator.print(f"Resumed from {checkpoint} at step {global_step}")
+
+    if accelerator.is_main_process:
+        summary = vars(args).copy()
+        summary["output_dir"] = str(args.output_dir)
+        summary["dataset_size"] = len(dataset)
+        summary["effective_batch_size"] = (
+            args.train_batch_size
+            * args.gradient_accumulation_steps
+            * accelerator.num_processes
+        )
+        (args.output_dir / "training-summary.json").write_text(
+            json.dumps(summary, indent=2, default=str) + "\n"
+        )
+
     data_iterator = iter(dataloader)
     while global_step < args.max_train_steps:
         try:
@@ -281,7 +433,7 @@ def main() -> None:
             data_iterator = iter(dataloader)
             batch = next(data_iterator)
 
-        with accelerator.accumulate(unet):
+        with accelerator.accumulate(student):
             pixel_values = batch["pixel_values"].to(
                 accelerator.device, dtype=weight_dtype, non_blocking=True
             )
@@ -296,7 +448,9 @@ def main() -> None:
                 latents = latents * vae.config.scaling_factor
                 masked_latents = vae.encode(masked_pixel_values).latent_dist.sample()
                 masked_latents = masked_latents * vae.config.scaling_factor
-                encoder_hidden_states = text_encoder(batch["input_ids"].to(accelerator.device))[0]
+                encoder_hidden_states = text_encoder(
+                    batch["input_ids"].to(accelerator.device)
+                )[0]
 
             noise = torch.randn_like(latents)
             timesteps = torch.randint(
@@ -306,61 +460,117 @@ def main() -> None:
                 device=latents.device,
             ).long()
             noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
-            mask = F.interpolate(mask, size=latents.shape[-2:], mode="nearest")
-            model_input = torch.cat([noisy_latents, mask, masked_latents], dim=1)
-            model_pred = unet(
+            latent_mask = F.interpolate(mask, size=latents.shape[-2:], mode="nearest")
+            model_input = torch.cat([noisy_latents, latent_mask, masked_latents], dim=1)
+            model_pred = student(
                 model_input,
                 timesteps,
                 encoder_hidden_states=encoder_hidden_states,
             ).sample
 
-            if noise_scheduler.config.prediction_type == "epsilon":
-                target = noise
-            elif noise_scheduler.config.prediction_type == "v_prediction":
-                target = noise_scheduler.get_velocity(latents, noise, timesteps)
+            prediction_type = noise_scheduler.config.prediction_type
+            if prediction_type == "epsilon":
+                ground_truth_target = noise
+            elif prediction_type == "v_prediction":
+                ground_truth_target = noise_scheduler.get_velocity(latents, noise, timesteps)
             else:
-                raise ValueError(
-                    f"Unsupported prediction type: {noise_scheduler.config.prediction_type}"
+                raise ValueError(f"Unsupported prediction type: {prediction_type}")
+
+            spatial_weights = spatial_loss_weights(
+                latent_mask,
+                context_weight=args.context_loss_weight,
+                masked_weight=args.masked_loss_weight,
+                boundary_weight=args.boundary_loss_weight,
+                boundary_radius=args.boundary_radius,
+            )
+            sample_weights = min_snr_weights(
+                noise_scheduler.alphas_cumprod,
+                timesteps,
+                gamma=args.snr_gamma,
+                prediction_type=prediction_type,
+            )
+            ground_truth_loss = weighted_mse(
+                model_pred,
+                ground_truth_target,
+                spatial_weights=spatial_weights,
+                sample_weights=sample_weights,
+            )
+            teacher_loss = torch.zeros((), device=accelerator.device)
+            if teacher is not None:
+                with torch.no_grad():
+                    teacher_target = teacher(
+                        model_input,
+                        timesteps,
+                        encoder_hidden_states=encoder_hidden_states,
+                    ).sample
+                teacher_loss = weighted_mse(
+                    model_pred,
+                    teacher_target,
+                    spatial_weights=spatial_weights,
+                    sample_weights=sample_weights,
                 )
-            loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+            loss = (
+                args.ground_truth_loss_weight * ground_truth_loss
+                + args.teacher_loss_weight * teacher_loss
+            )
+
             accelerator.backward(loss)
             if accelerator.sync_gradients:
-                accelerator.clip_grad_norm_(unet.parameters(), 1.0)
+                accelerator.clip_grad_norm_(student.parameters(), 1.0)
             optimizer.step()
             lr_scheduler.step()
             optimizer.zero_grad(set_to_none=True)
 
         if accelerator.sync_gradients:
             global_step += 1
-            if accelerator.is_main_process and global_step % 50 == 0:
+            metrics = {
+                "step": global_step,
+                "loss": loss.detach().float().item(),
+                "teacher_loss": teacher_loss.detach().float().item(),
+                "ground_truth_loss": ground_truth_loss.detach().float().item(),
+                "lr": lr_scheduler.get_last_lr()[0],
+                "mask_area": batch["mask_area"].float().mean().item(),
+            }
+            if accelerator.is_main_process and (global_step == 1 or global_step % 25 == 0):
                 accelerator.print(
                     f"step={global_step}/{args.max_train_steps} "
-                    f"loss={loss.detach().item():.4f} "
-                    f"lr={lr_scheduler.get_last_lr()[0]:.3e}"
+                    f"loss={metrics['loss']:.5f} "
+                    f"teacher={metrics['teacher_loss']:.5f} "
+                    f"ground_truth={metrics['ground_truth_loss']:.5f} "
+                    f"mask={metrics['mask_area']:.3f} "
+                    f"lr={metrics['lr']:.3e}"
                 )
-            accelerator.log(
-                {"train_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]},
-                step=global_step,
-            )
+                with metrics_path.open("a") as handle:
+                    handle.write(json.dumps(metrics) + "\n")
+
+            if global_step % args.checkpointing_steps == 0:
+                accelerator.wait_for_everyone()
+                checkpoint = checkpoint_root / f"checkpoint-{global_step}"
+                accelerator.save_state(checkpoint, safe_serialization=True)
+                if accelerator.is_main_process:
+                    _prune_checkpoints(checkpoint_root, args.checkpoints_total_limit)
+                    (args.output_dir / "progress.json").write_text(
+                        json.dumps(metrics, indent=2) + "\n"
+                    )
+                accelerator.print(f"Saved resumable checkpoint {checkpoint}")
 
     accelerator.wait_for_everyone()
     if accelerator.is_main_process:
-        unwrapped = accelerator.unwrap_model(unet).cpu()
-        save_pipeline(
-            base_model=args.pretrained_model_name_or_path,
-            revision=args.revision,
-            unet=unwrapped,
-            output_dir=args.output_dir,
-        )
-        summary = vars(args).copy()
-        summary["output_dir"] = str(args.output_dir)
-        (args.output_dir / "training-summary.json").write_text(
-            json.dumps(summary, indent=2, default=str) + "\n"
+        unwrapped = accelerator.unwrap_model(student).cpu()
+        _save_pipeline(args=args, unet=unwrapped, output_dir=args.output_dir)
+        final_metrics = {
+            "completed_steps": global_step,
+            "final_loss": loss.detach().float().item(),
+            "final_teacher_loss": teacher_loss.detach().float().item(),
+            "final_ground_truth_loss": ground_truth_loss.detach().float().item(),
+        }
+        (args.output_dir / "training-complete.json").write_text(
+            json.dumps(final_metrics, indent=2) + "\n"
         )
         if args.push_to_hub:
             pipeline = StableDiffusionInpaintPipeline.from_pretrained(args.output_dir)
             pipeline.push_to_hub(args.hub_model_id)
-        accelerator.print(f"Saved Clover Image Tiny Inpaint to {args.output_dir}")
+        accelerator.print(f"Saved Clover Image Tiny Inpaint v2 to {args.output_dir}")
     accelerator.end_training()
 
 

modal_inpaint.py

@@ -1,24 +1,21 @@
 #!/usr/bin/env python3
-"""Run Clover Image Tiny inpainting training on Modal.
-
-The default job writes its result to a persistent Modal Volume. This keeps
-training independent from Hub credentials; the trained directory can be
-downloaded and uploaded to the model repository after validation.
-"""
+"""Run context-aware Clover Image Tiny inpainting distillation on Modal."""
 
 from __future__ import annotations
 
 import os
 import subprocess
 import sys
+import time
 from pathlib import Path
 
 import modal
 
-
-APP_NAME = "clover-image-tiny-inpaint"
+APP_NAME = "clover-image-tiny-inpaint-v2"
 OUTPUT_VOLUME_NAME = "clover-image-tiny-inpaint-output"
+CACHE_VOLUME_NAME = "clover-image-tiny-inpaint-cache"
 OUTPUT_ROOT = Path("/outputs")
+CACHE_ROOT = Path("/cache")
 
 image = (
     modal.Image.debian_slim(python_version="3.11")
@@ -39,42 +36,69 @@ image = (
 )
 
 output_volume = modal.Volume.from_name(OUTPUT_VOLUME_NAME, create_if_missing=True)
+cache_volume = modal.Volume.from_name(CACHE_VOLUME_NAME, create_if_missing=True)
 app = modal.App(
     APP_NAME,
     image=image,
-    volumes={str(OUTPUT_ROOT): output_volume},
+    volumes={
+        str(OUTPUT_ROOT): output_volume,
+        str(CACHE_ROOT): cache_volume,
+    },
 )
 
 
-@app.function(gpu="A10G", timeout=4 * 60 * 60)
+@app.function(gpu="A10", timeout=12 * 60 * 60, cpu=8, memory=32768)
 def train(
     *,
     base_model: str = "neonforestmist/Clover-Image-Tiny",
     base_revision: str = "63b0e9f6be9c00888ff464f342a9ef052bf76681",
+    initial_inpaint_model: str = "neonforestmist/Clover-Image-Tiny-Inpaint",
+    initial_inpaint_revision: str = "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
+    teacher_model: str = "stable-diffusion-v1-5/stable-diffusion-inpainting",
+    teacher_revision: str = "8a4288a76071f7280aedbdb3253bdb9e9d5d84bb",
     dataset: str = "prithivMLmods/Caption3o-Opt",
     dataset_revision: str = "17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b",
     dataset_split: str = "train",
     image_column: str = "image",
     caption_column: str = "caption",
-    max_train_steps: int = 4000,
-    output_name: str = "clover-image-tiny-inpaint",
+    max_train_steps: int = 12000,
+    output_name: str = "clover-image-tiny-inpaint-v2",
     max_train_samples: int | None = None,
-    learning_rate: float = 1e-5,
-    seed: int = 20260810,
+    learning_rate: float = 5e-6,
+    lr_warmup_steps: int | None = None,
+    teacher_loss_weight: float = 0.75,
+    ground_truth_loss_weight: float = 0.25,
+    context_loss_weight: float = 0.25,
+    masked_loss_weight: float = 2.5,
+    boundary_loss_weight: float = 2.0,
+    seed: int = 20260811,
+    checkpointing_steps: int = 500,
+    resume: bool = False,
 ) -> str:
-    """Train once and return the Modal Volume path containing the pipeline."""
+    """Train v2 and return its persistent Modal Volume path."""
 
     output_dir = OUTPUT_ROOT / output_name
-    if output_dir.exists():
-        raise RuntimeError(f"Output already exists; choose another --output-name: {output_dir}")
+    if (output_dir / "training-complete.json").exists():
+        raise RuntimeError(f"Completed output already exists: {output_dir}")
+    if output_dir.exists() and not resume:
+        raise RuntimeError(
+            f"Partial output already exists; choose another name or enable resume: {output_dir}"
+        )
 
     command = [
         sys.executable,
+        "-u",
         "/root/inpainting/train.py",
         "--pretrained_model_name_or_path",
         base_model,
         "--revision",
         base_revision,
+        "--initial_inpaint_model",
+        initial_inpaint_model,
+        "--teacher_model_name_or_path",
+        teacher_model,
+        "--teacher_revision",
+        teacher_revision,
         "--dataset_name",
         dataset,
         "--dataset_revision",
@@ -89,43 +113,118 @@ def train(
         str(max_train_steps),
         "--learning_rate",
         str(learning_rate),
+        "--lr_warmup_steps",
+        str(
+            lr_warmup_steps
+            if lr_warmup_steps is not None
+            else min(500, max(1, max_train_steps // 20))
+        ),
         "--gradient_accumulation_steps",
         "4",
         "--gradient_checkpointing",
         "--mixed_precision",
-        "fp16",
+        "bf16",
+        "--random_flip",
+        "--caption_dropout_probability",
+        "0.10",
+        "--validation_samples",
+        "128",
+        "--teacher_loss_weight",
+        str(teacher_loss_weight),
+        "--ground_truth_loss_weight",
+        str(ground_truth_loss_weight),
+        "--context_loss_weight",
+        str(context_loss_weight),
+        "--masked_loss_weight",
+        str(masked_loss_weight),
+        "--boundary_loss_weight",
+        str(boundary_loss_weight),
+        "--snr_gamma",
+        "5.0",
+        "--checkpointing_steps",
+        str(checkpointing_steps),
+        "--checkpoints_total_limit",
+        "3",
         "--seed",
         str(seed),
         "--output_dir",
         str(output_dir),
     ]
+    if initial_inpaint_revision:
+        command.extend(["--initial_inpaint_revision", initial_inpaint_revision])
     if max_train_samples is not None:
         command.extend(["--max_train_samples", str(max_train_samples)])
+    if resume:
+        command.extend(["--resume_from_checkpoint", "latest"])
 
     env = os.environ.copy()
-    env.setdefault("HF_HOME", "/root/.cache/huggingface")
-    subprocess.run(command, check=True, env=env)
+    env.update(
+        {
+            "HF_HOME": str(CACHE_ROOT / "huggingface"),
+            "HF_HUB_CACHE": str(CACHE_ROOT / "huggingface" / "hub"),
+            "HF_DATASETS_CACHE": str(CACHE_ROOT / "huggingface" / "datasets"),
+            "TOKENIZERS_PARALLELISM": "false",
+            "PYTHONUNBUFFERED": "1",
+        }
+    )
+    process = subprocess.Popen(command, env=env)
+    last_commit = time.monotonic()
+    while process.poll() is None:
+        time.sleep(30)
+        if time.monotonic() - last_commit >= 300:
+            output_volume.commit()
+            cache_volume.commit()
+            last_commit = time.monotonic()
+    if process.returncode:
+        output_volume.commit()
+        cache_volume.commit()
+        raise subprocess.CalledProcessError(process.returncode, command)
+
     output_volume.commit()
+    cache_volume.commit()
     return str(output_dir)
 
 
 @app.local_entrypoint()
 def main(
     smoke: bool = False,
-    steps: int = 4000,
-    output_name: str = "clover-image-tiny-inpaint",
+    steps: int = 12000,
+    output_name: str = "clover-image-tiny-inpaint-v2",
+    initial_inpaint_model: str = "neonforestmist/Clover-Image-Tiny-Inpaint",
+    initial_inpaint_revision: str | None = "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
+    learning_rate: float = 5e-6,
+    lr_warmup_steps: int | None = None,
+    teacher_loss_weight: float = 0.75,
+    ground_truth_loss_weight: float = 0.25,
+    context_loss_weight: float = 0.25,
+    masked_loss_weight: float = 2.5,
+    boundary_loss_weight: float = 2.0,
+    resume: bool = False,
 ) -> None:
-    """Launch a bounded smoke job or the full A10G run."""
+    """Launch a small pipeline test or the full bounded A100 run."""
 
     if smoke:
-        steps = min(steps, 2)
-        samples = 4
-        output_name = f"{output_name}-smoke"
+        steps = min(steps, 4)
+        samples = 8
+        output_name = f"{output_name}-smoke-{int(time.time())}"
+        checkpointing_steps = 2
     else:
         samples = None
+        checkpointing_steps = 500
     result = train.remote(
         max_train_steps=steps,
         max_train_samples=samples,
         output_name=output_name,
+        initial_inpaint_model=initial_inpaint_model,
+        initial_inpaint_revision=initial_inpaint_revision,
+        learning_rate=learning_rate,
+        lr_warmup_steps=lr_warmup_steps,
+        teacher_loss_weight=teacher_loss_weight,
+        ground_truth_loss_weight=ground_truth_loss_weight,
+        context_loss_weight=context_loss_weight,
+        masked_loss_weight=masked_loss_weight,
+        boundary_loss_weight=boundary_loss_weight,
+        checkpointing_steps=checkpointing_steps,
+        resume=resume,
     )
     print(f"Training output is available in Modal Volume {OUTPUT_VOLUME_NAME}: {result}")

modal_inpaint_eval.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+"""Evaluate a Modal Volume Clover inpainting candidate against the release."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import modal
+
+APP_NAME = "clover-image-tiny-inpaint-eval"
+OUTPUT_VOLUME_NAME = "clover-image-tiny-inpaint-output"
+CACHE_VOLUME_NAME = "clover-image-tiny-inpaint-cache"
+OUTPUT_ROOT = Path("/outputs")
+CACHE_ROOT = Path("/cache")
+
+image = (
+    modal.Image.debian_slim(python_version="3.11")
+    .pip_install(
+        "accelerate==1.14.0",
+        "datasets==4.8.5",
+        "diffusers==0.39.0",
+        "huggingface_hub==0.36.0",
+        "numpy==2.2.6",
+        "pillow==12.3.0",
+        "safetensors==0.8.0",
+        "torch==2.7.0",
+        "torchvision==0.22.0",
+        "transformers==4.57.6",
+    )
+    .add_local_dir("inpainting", remote_path="/root/inpainting")
+)
+
+output_volume = modal.Volume.from_name(OUTPUT_VOLUME_NAME, create_if_missing=True)
+cache_volume = modal.Volume.from_name(CACHE_VOLUME_NAME, create_if_missing=True)
+app = modal.App(
+    APP_NAME,
+    image=image,
+    volumes={
+        str(OUTPUT_ROOT): output_volume,
+        str(CACHE_ROOT): cache_volume,
+    },
+)
+
+
+@app.function(gpu="A10", timeout=2 * 60 * 60, cpu=4, memory=24576)
+def evaluate(candidate_name: str, evaluation_name: str, sample_count: int = 6) -> str:
+    candidate = OUTPUT_ROOT / candidate_name
+    if not (candidate / "training-complete.json").exists():
+        raise RuntimeError(f"Candidate is not complete: {candidate}")
+    destination = OUTPUT_ROOT / "evaluations" / evaluation_name
+    if destination.exists():
+        raise RuntimeError(f"Evaluation output already exists: {destination}")
+    env = os.environ.copy()
+    env.update(
+        {
+            "HF_HOME": str(CACHE_ROOT / "huggingface"),
+            "HF_HUB_CACHE": str(CACHE_ROOT / "huggingface" / "hub"),
+            "HF_DATASETS_CACHE": str(CACHE_ROOT / "huggingface" / "datasets"),
+            "TOKENIZERS_PARALLELISM": "false",
+        }
+    )
+    command = [
+        sys.executable,
+        "-u",
+        "/root/inpainting/evaluate.py",
+        "--baseline_model",
+        "neonforestmist/Clover-Image-Tiny-Inpaint",
+        "--baseline_revision",
+        "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
+        "--candidate_model",
+        str(candidate),
+        "--dataset_name",
+        "prithivMLmods/Caption3o-Opt",
+        "--dataset_revision",
+        "17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b",
+        "--sample_count",
+        str(sample_count),
+        "--output_dir",
+        str(destination),
+    ]
+    subprocess.run(command, check=True, env=env)
+    output_volume.commit()
+    cache_volume.commit()
+    return str(destination)
+
+
+@app.local_entrypoint()
+def main(candidate_name: str, evaluation_name: str, sample_count: int = 6) -> None:
+    result = evaluate.remote(candidate_name, evaluation_name, sample_count)
+    print(f"Evaluation is available in Modal Volume {OUTPUT_VOLUME_NAME}: {result}")

modal_inpaint_semantic_eval.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+"""Run controlled semantic inpainting evaluation from a Modal Volume candidate."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import modal
+
+APP_NAME = "clover-image-tiny-inpaint-semantic-eval"
+OUTPUT_VOLUME_NAME = "clover-image-tiny-inpaint-output"
+CACHE_VOLUME_NAME = "clover-image-tiny-inpaint-cache"
+OUTPUT_ROOT = Path("/outputs")
+CACHE_ROOT = Path("/cache")
+
+image = (
+    modal.Image.debian_slim(python_version="3.11")
+    .pip_install(
+        "accelerate==1.14.0",
+        "diffusers==0.39.0",
+        "numpy==2.2.6",
+        "pillow==12.3.0",
+        "safetensors==0.8.0",
+        "torch==2.7.0",
+        "torchvision==0.22.0",
+        "transformers==4.57.6",
+    )
+    .add_local_dir("inpainting", remote_path="/root/inpainting")
+)
+
+output_volume = modal.Volume.from_name(OUTPUT_VOLUME_NAME, create_if_missing=True)
+cache_volume = modal.Volume.from_name(CACHE_VOLUME_NAME, create_if_missing=True)
+app = modal.App(
+    APP_NAME,
+    image=image,
+    volumes={
+        str(OUTPUT_ROOT): output_volume,
+        str(CACHE_ROOT): cache_volume,
+    },
+)
+
+
+@app.function(gpu="A10", timeout=3 * 60 * 60, cpu=4, memory=24576)
+def evaluate(
+    candidate_name: str,
+    evaluation_name: str,
+    guidance_scale: float = 7.5,
+    steps: int = 30,
+    mask_crop_padding: int = 0,
+) -> str:
+    candidate = OUTPUT_ROOT / candidate_name
+    if not (candidate / "training-complete.json").exists():
+        raise RuntimeError(f"Candidate is not complete: {candidate}")
+    destination = OUTPUT_ROOT / "evaluations" / evaluation_name
+    if destination.exists():
+        raise RuntimeError(f"Evaluation output already exists: {destination}")
+    env = os.environ.copy()
+    env.update(
+        {
+            "HF_HOME": str(CACHE_ROOT / "huggingface"),
+            "HF_HUB_CACHE": str(CACHE_ROOT / "huggingface" / "hub"),
+            "TOKENIZERS_PARALLELISM": "false",
+        }
+    )
+    command = [
+        sys.executable,
+        "-u",
+        "/root/inpainting/evaluate_semantic.py",
+        "--base_model",
+        "neonforestmist/Clover-Image-Tiny",
+        "--base_revision",
+        "63b0e9f6be9c00888ff464f342a9ef052bf76681",
+        "--baseline_model",
+        "neonforestmist/Clover-Image-Tiny-Inpaint",
+        "--baseline_revision",
+        "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
+        "--teacher_model",
+        "stable-diffusion-v1-5/stable-diffusion-inpainting",
+        "--teacher_revision",
+        "8a4288a76071f7280aedbdb3253bdb9e9d5d84bb",
+        "--teacher_variant",
+        "fp16",
+        "--candidate_model",
+        str(candidate),
+        "--clip_model",
+        "openai/clip-vit-base-patch32",
+        "--clip_revision",
+        "3d74acf9a28c67741b2f4f2ea7635f0aaf6f0268",
+        "--output_dir",
+        str(destination),
+        "--guidance_scale",
+        str(guidance_scale),
+        "--steps",
+        str(steps),
+        "--mask_crop_padding",
+        str(mask_crop_padding),
+    ]
+    subprocess.run(command, check=True, env=env)
+    output_volume.commit()
+    cache_volume.commit()
+    return str(destination)
+
+
+@app.local_entrypoint()
+def main(
+    candidate_name: str,
+    evaluation_name: str,
+    guidance_scale: float = 7.5,
+    steps: int = 30,
+    mask_crop_padding: int = 0,
+) -> None:
+    result = evaluate.remote(
+        candidate_name,
+        evaluation_name,
+        guidance_scale,
+        steps,
+        mask_crop_padding,
+    )
+    print(f"Evaluation is available in Modal Volume {OUTPUT_VOLUME_NAME}: {result}")

modal_inpaint_snapshot.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+"""Materialize a resumable training checkpoint as a Diffusers pipeline."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import modal
+
+APP_NAME = "clover-image-tiny-inpaint-snapshot"
+OUTPUT_VOLUME_NAME = "clover-image-tiny-inpaint-output"
+CACHE_VOLUME_NAME = "clover-image-tiny-inpaint-cache"
+OUTPUT_ROOT = Path("/outputs")
+CACHE_ROOT = Path("/cache")
+INITIAL_MODEL = "neonforestmist/Clover-Image-Tiny-Inpaint"
+INITIAL_REVISION = "1b6f8ae3db51900520369d5522c7dc7c2a97e21e"
+
+image = modal.Image.debian_slim(python_version="3.11").pip_install(
+    "accelerate==1.14.0",
+    "diffusers==0.39.0",
+    "huggingface_hub==0.36.0",
+    "safetensors==0.8.0",
+    "torch==2.7.0",
+    "transformers==4.57.6",
+)
+
+output_volume = modal.Volume.from_name(OUTPUT_VOLUME_NAME, create_if_missing=True)
+cache_volume = modal.Volume.from_name(CACHE_VOLUME_NAME, create_if_missing=True)
+app = modal.App(
+    APP_NAME,
+    image=image,
+    volumes={
+        str(OUTPUT_ROOT): output_volume,
+        str(CACHE_ROOT): cache_volume,
+    },
+)
+
+
+@app.function(timeout=60 * 60, cpu=4, memory=16384)
+def materialize(source_name: str, checkpoint_step: int, output_name: str) -> str:
+    import torch
+    from diffusers import StableDiffusionInpaintPipeline
+    from safetensors.torch import load_file
+
+    source = OUTPUT_ROOT / source_name
+    checkpoint = source / "checkpoints" / f"checkpoint-{checkpoint_step}"
+    weights = checkpoint / "model.safetensors"
+    if not weights.is_file():
+        raise RuntimeError(f"Missing checkpoint weights: {weights}")
+    destination = OUTPUT_ROOT / output_name
+    if destination.exists():
+        raise RuntimeError(f"Snapshot output already exists: {destination}")
+
+    os.environ.update(
+        {
+            "HF_HOME": str(CACHE_ROOT / "huggingface"),
+            "HF_HUB_CACHE": str(CACHE_ROOT / "huggingface" / "hub"),
+            "TOKENIZERS_PARALLELISM": "false",
+        }
+    )
+    pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+        INITIAL_MODEL,
+        revision=INITIAL_REVISION,
+        torch_dtype=torch.float32,
+    )
+    state = load_file(str(weights), device="cpu")
+    incompatible = pipeline.unet.load_state_dict(state, strict=True)
+    if incompatible.missing_keys or incompatible.unexpected_keys:
+        raise RuntimeError(f"Checkpoint state mismatch: {incompatible}")
+    pipeline.save_pretrained(destination, safe_serialization=True)
+
+    metadata = {
+        "snapshot_type": "training_checkpoint",
+        "source_name": source_name,
+        "checkpoint_step": checkpoint_step,
+        "initial_model": INITIAL_MODEL,
+        "initial_revision": INITIAL_REVISION,
+    }
+    (destination / "snapshot.json").write_text(json.dumps(metadata, indent=2) + "\n")
+    (destination / "training-complete.json").write_text(
+        json.dumps(metadata, indent=2) + "\n"
+    )
+    output_volume.commit()
+    cache_volume.commit()
+    return str(destination)
+
+
+@app.local_entrypoint()
+def main(source_name: str, checkpoint_step: int, output_name: str) -> None:
+    result = materialize.remote(source_name, checkpoint_step, output_name)
+    print(f"Snapshot is available in Modal Volume {OUTPUT_VOLUME_NAME}: {result}")

training-complete.json

@@ -0,0 +1,6 @@
+{
+  "completed_steps": 500,
+  "final_loss": 0.11363611370325089,
+  "final_teacher_loss": 0.05577564239501953,
+  "final_ground_truth_loss": 0.28721752762794495
+}

training-summary.json

@@ -1,6 +1,10 @@
 {
   "pretrained_model_name_or_path": "neonforestmist/Clover-Image-Tiny",
   "revision": "63b0e9f6be9c00888ff464f342a9ef052bf76681",
+  "initial_inpaint_model": "neonforestmist/Clover-Image-Tiny-Inpaint",
+  "initial_inpaint_revision": "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
+  "teacher_model_name_or_path": "stable-diffusion-v1-5/stable-diffusion-inpainting",
+  "teacher_revision": "8a4288a76071f7280aedbdb3253bdb9e9d5d84bb",
   "dataset_name": "prithivMLmods/Caption3o-Opt",
   "dataset_config_name": null,
   "dataset_revision": "17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b",
@@ -9,19 +13,34 @@
   "caption_column": "caption",
   "resolution": 512,
   "train_batch_size": 1,
-  "max_train_steps": 4000,
-  "learning_rate": 1e-05,
+  "max_train_steps": 500,
+  "learning_rate": 5e-06,
   "lr_scheduler": "cosine",
-  "lr_warmup_steps": 200,
+  "lr_warmup_steps": 25,
   "gradient_accumulation_steps": 4,
   "gradient_checkpointing": true,
-  "mixed_precision": "fp16",
-  "seed": 20260810,
-  "mask_min_area": 0.12,
-  "mask_max_area": 0.55,
+  "mixed_precision": "bf16",
+  "seed": 20260811,
+  "mask_min_area": 0.04,
+  "mask_max_area": 0.65,
+  "caption_dropout_probability": 0.1,
+  "teacher_loss_weight": 0.75,
+  "ground_truth_loss_weight": 0.25,
+  "context_loss_weight": 0.25,
+  "masked_loss_weight": 2.5,
+  "boundary_loss_weight": 2.0,
+  "boundary_radius": 2,
+  "snr_gamma": 5.0,
+  "random_flip": true,
   "max_train_samples": null,
-  "output_dir": "/outputs/clover-image-tiny-inpaint",
-  "validation_prompt": null,
+  "validation_samples": 128,
+  "num_workers": 4,
+  "checkpointing_steps": 500,
+  "checkpoints_total_limit": 3,
+  "resume_from_checkpoint": null,
+  "output_dir": "/outputs/clover-image-tiny-inpaint-v2-pilot-500",
   "push_to_hub": false,
-  "hub_model_id": null
+  "hub_model_id": null,
+  "dataset_size": 10150,
+  "effective_batch_size": 4
 }

training/README-INPAINTING.md

@@ -9,10 +9,20 @@ The U-Net input is ordered as:
 [noisy_latent (4), mask (1), masked_image_latent (4)]
 ```
 
-White mask pixels are regenerated; black pixels are preserved. Masks are
-sampled procedurally during training, so the dataset only needs image/caption
-pairs. The pinned default dataset is `prithivMLmods/Caption3o-Opt` at revision
-`17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b` (Apache-2.0, `image` + `caption`).
+White mask pixels are regenerated; black pixels are preserved. The v2 recipe
+warm-starts from the released nine-channel checkpoint and distills the pinned
+official `stable-diffusion-v1-5/stable-diffusion-inpainting` U-Net. A smaller
+ground-truth diffusion term prevents the student from merely copying teacher
+errors. Min-SNR timestep weighting and a spatial objective emphasize the
+masked region, its boundary, and enough unmasked context to learn coherent
+transitions.
+
+Masks are sampled procedurally as free-form strokes, multiple disconnected
+strokes, rounded boxes, ellipses, polygons, multiple object-like regions, and
+outpainting bands. The pinned dataset is `prithivMLmods/Caption3o-Opt` at
+revision `17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b` (Apache-2.0, `image` +
+`caption`). A deterministic 128-image tail is excluded from training and used
+for same-mask, same-seed checkpoint comparison.
 
 ## Modal
 
@@ -21,26 +31,35 @@ Use that profile before running the commands below:
 
 ```bash
 modal profile activate guccichungus69
-modal run modal_inpaint.py --smoke --steps 2
-modal run modal_inpaint.py --steps 4000
+modal run modal_inpaint.py --smoke --steps 1
+modal run --detach modal_inpaint.py \
+  --steps 500 \
+  --lr-warmup-steps 25 \
+  --output-name clover-image-tiny-inpaint-v2-pilot-500
 ```
 
-The job uses one A10G and writes the finished Diffusers pipeline to the
+The job uses one A10 and writes the finished Diffusers pipeline to the
 `clover-image-tiny-inpaint-output` Volume. No Hub token is required for the
-default public dataset. The estimated compute ceiling for the default four-hour
-job is about $4 at $1/hour; the actual charge depends on runtime.
+default public dataset. Resumable Accelerator checkpoints are committed to the
+persistent Volume. A separate cache avoids downloading the student, teacher,
+and dataset again on every run.
+
+The published v2 checkpoint is the 500-step cosine-decay sweep. A longer run
+was rejected after its held-out reconstruction and semantic gates regressed
+into colored high-frequency artifacts. The release pins the best validated
+checkpoint rather than the final checkpoint from the largest job.
 
 Download a completed output directory with:
 
 ```bash
 modal volume get clover-image-tiny-inpaint-output \
-  clover-image-tiny-inpaint \
-  ./artifacts/clover-image-tiny-inpaint
+  clover-image-tiny-inpaint-v2-pilot-500/unet \
+  ./artifacts/clover-image-tiny-inpaint-v2
 ```
 
-If a job stops before writing its final pipeline, the Volume retains any files
-that were committed. The current trainer saves only at the end, so use a
-separate `--output-name` for retries.
+If a job stops after a committed checkpoint, resume the same output with
+`--resume`. Completed output directories are immutable to the launcher; use a
+new `--output-name` for a distinct experiment.
 
 ## Local dry run
 

unet/config.json

@@ -1,7 +1,7 @@
 {
   "_class_name": "UNet2DConditionModel",
   "_diffusers_version": "0.39.0",
-  "_name_or_path": "neonforestmist/Clover-Image-Tiny",
+  "_name_or_path": "neonforestmist/Clover-Image-Tiny-Inpaint",
   "act_fn": "silu",
   "addition_embed_type": null,
   "addition_embed_type_num_heads": 64,