Add Modal training and inpainting model scaffold

inpainting/README.md

@@ -0,0 +1,68 @@
+---
+library_name: diffusers
+pipeline_tag: image-inpainting
+base_model: neonforestmist/Clover-Image-Tiny
+license: creativeml-openrail-m
+tags:
+  - clover-image
+  - 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.
+
+The model uses a 9-channel U-Net input:
+
+```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 iPhone.
+
+## Diffusers
+
+```python
+from diffusers import AutoPipelineForInpainting
+from diffusers.utils import load_image
+
+pipe = AutoPipelineForInpainting.from_pretrained(
+    "neonforestmist/Clover-Image-Tiny-Inpaint",
+    torch_dtype="auto",
+)
+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,
+).images[0]
+image.save("clover-inpaint.png")
+```
+
+## Core ML and iPhone 15
+
+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
+bundled `VAEEncoder.mlmodelc` creates the masked-image latent locally, so the
+input image and mask do not leave the device.
+
+Conversion and the native iOS integration live in the source Clover repo:
+
+- [`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)
+
+## 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.

inpainting/__init__.py

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

inpainting/config.json

@@ -0,0 +1,29 @@
+{
+  "base_model": "neonforestmist/Clover-Image-Tiny",
+  "base_revision": "63b0e9f6be9c00888ff464f342a9ef052bf76681",
+  "dataset": "prithivMLmods/Caption3o-Opt",
+  "dataset_revision": "17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b",
+  "dataset_split": "train",
+  "image_column": "image",
+  "caption_column": "caption",
+  "resolution": 512,
+  "unet_in_channels": 9,
+  "mask_semantics": "white=regenerate, black=preserve",
+  "training": {
+    "steps": 4000,
+    "learning_rate": 0.00001,
+    "warmup_steps": 200,
+    "train_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "mixed_precision": "fp16",
+    "gradient_checkpointing": true,
+    "mask_min_area": 0.12,
+    "mask_max_area": 0.55
+  },
+  "modal": {
+    "profile": "guccichungus69",
+    "gpu": "A10G",
+    "timeout_hours": 4,
+    "estimated_hourly_usd": 1.0
+  }
+}

inpainting/masks.py

@@ -0,0 +1,62 @@
+"""Synthetic mask generation used for inpainting fine-tuning."""
+
+from __future__ import annotations
+
+import random
+
+from PIL import Image, ImageChops, ImageDraw
+
+
+def random_mask(
+    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"."""
+
+    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),
+            )
+            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)
+
+    return mask
+
+
+def apply_mask(image: Image.Image, mask: Image.Image) -> Image.Image:
+    """Black out the pixels that the inpainting model must regenerate."""
+
+    image = image.convert("RGB")
+    keep = ImageChops.invert(mask.convert("L"))
+    return Image.composite(image, Image.new("RGB", image.size), keep)

inpainting/model.py

@@ -0,0 +1,57 @@
+"""Model construction helpers for Clover Image Tiny inpainting."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+from diffusers import UNet2DConditionModel
+
+
+def _load_kwargs(revision: str | None) -> dict[str, Any]:
+    return {"revision": revision} if revision else {}
+
+
+def make_inpainting_unet(
+    base_model: str | Path,
+    *,
+    revision: str | None = None,
+    zero_initialize_conditioning: bool = True,
+) -> UNet2DConditionModel:
+    """Create a 9-channel U-Net from the 4-channel Clover checkpoint.
+
+    The first four input channels retain Clover's original weights. The extra
+    channels receive the binary inpainting mask and the masked-image latent.
+    Zero-initializing them preserves a stable text-to-image starting point while
+    the inpainting fine-tune learns how to use the new conditioning channels.
+    """
+
+    load_kwargs = _load_kwargs(revision)
+    base = UNet2DConditionModel.from_pretrained(
+        str(base_model),
+        subfolder="unet",
+        low_cpu_mem_usage=False,
+        **load_kwargs,
+    )
+    config = dict(base.config)
+    config["in_channels"] = 9
+    inpaint = UNet2DConditionModel.from_config(config)
+
+    state = base.state_dict()
+    conv_in_weight = state.pop("conv_in.weight")
+    inpaint.load_state_dict(state, strict=False)
+
+    with torch.no_grad():
+        if zero_initialize_conditioning:
+            inpaint.conv_in.weight.zero_()
+            inpaint.conv_in.weight[:, :4].copy_(conv_in_weight)
+        else:
+            # Keep the pretrained channels and leave the five new channels at
+            # the framework's default initialization.
+            inpaint.conv_in.weight[:, :4].copy_(conv_in_weight)
+        if "conv_in.bias" in base.state_dict():
+            inpaint.conv_in.bias.copy_(base.conv_in.bias)
+
+    del base
+    return inpaint

inpainting/train.py

@@ -0,0 +1,368 @@
+#!/usr/bin/env python3
+"""Fine-tune Clover Image Tiny into a 9-channel inpainting pipeline.
+
+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.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import shutil
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from accelerate import Accelerator
+from datasets import load_dataset
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    StableDiffusionInpaintPipeline,
+    get_scheduler,
+)
+from PIL import Image
+from transformers import CLIPTextModel, CLIPTokenizer
+from torchvision import transforms
+
+from inpainting.masks import apply_mask, random_mask
+from inpainting.model import make_inpainting_unet
+
+
+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("--dataset_name", required=True)
+    parser.add_argument("--dataset_config_name")
+    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("--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("--lr_scheduler", default="cosine")
+    parser.add_argument("--lr_warmup_steps", type=int, default=200)
+    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("--max_train_samples", type=int)
+    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()
+
+
+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 make_collate_fn(
+    *,
+    tokenizer: CLIPTokenizer,
+    resolution: int,
+    min_area: float,
+    max_area: float,
+    seed: int,
+):
+    worker_rng = random.Random(seed)
+
+    def collate(examples: list[dict[str, Any]]) -> dict[str, Any]:
+        images: list[torch.Tensor] = []
+        masked_images: list[torch.Tensor] = []
+        masks: list[torch.Tensor] = []
+        captions: list[str] = []
+
+        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)
+            mask = random_mask(
+                (resolution, resolution),
+                worker_rng,
+                min_area=min_area,
+                max_area=max_area,
+            )
+            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 ""))
+
+        tokenized = tokenizer(
+            captions,
+            max_length=tokenizer.model_max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+        return {
+            "pixel_values": torch.stack(images),
+            "masked_pixel_values": torch.stack(masked_images),
+            "mask": torch.stack(masks),
+            "input_ids": tokenized.input_ids,
+        }
+
+    return collate
+
+
+def save_pipeline(
+    *,
+    base_model: str,
+    revision: str | None,
+    unet,
+    output_dir: Path,
+) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+        base_model,
+        revision=revision,
+        unet=unet,
+    )
+    pipeline.save_pretrained(output_dir, safe_serialization=True)
+    (output_dir / "inpainting-config.json").write_text(
+        json.dumps(
+            {
+                "unet_in_channels": 9,
+                "mask_semantics": "white=regenerate, black=preserve",
+                "base_model": base_model,
+                "base_revision": revision,
+            },
+            indent=2,
+        )
+        + "\n"
+    )
+
+
+def main() -> None:
+    args = parse_args()
+    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 args.push_to_hub and not args.hub_model_id:
+        raise ValueError("--hub_model_id is required with --push_to_hub")
+
+    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)
+
+    kwargs = model_kwargs(args.revision)
+    dataset = load_dataset(
+        args.dataset_name,
+        args.dataset_config_name,
+        split=args.dataset_split,
+        revision=args.dataset_revision,
+    )
+    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:
+        raise ValueError(f"Missing image column {args.image_column!r}: {dataset.column_names}")
+    if args.caption_column not in dataset.column_names:
+        raise ValueError(f"Missing caption column {args.caption_column!r}: {dataset.column_names}")
+    dataset = dataset.rename_columns(
+        {args.image_column: "image", args.caption_column: "caption"}
+    )
+
+    tokenizer = CLIPTokenizer.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="tokenizer",
+        **kwargs,
+    )
+    text_encoder = CLIPTextModel.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="text_encoder",
+        **kwargs,
+    )
+    vae = AutoencoderKL.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="vae",
+        **kwargs,
+    )
+    unet = make_inpainting_unet(
+        args.pretrained_model_name_or_path,
+        revision=args.revision,
+    )
+    noise_scheduler = DDPMScheduler.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="scheduler",
+        **kwargs,
+    )
+
+    if args.gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
+    text_encoder.requires_grad_(False)
+    vae.requires_grad_(False)
+    text_encoder.eval()
+    vae.eval()
+
+    collate_fn = make_collate_fn(
+        tokenizer=tokenizer,
+        resolution=args.resolution,
+        min_area=args.mask_min_area,
+        max_area=args.mask_max_area,
+        seed=args.seed,
+    )
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        shuffle=True,
+        collate_fn=collate_fn,
+        batch_size=args.train_batch_size,
+        num_workers=2,
+        pin_memory=True,
+    )
+    optimizer = torch.optim.AdamW(
+        unet.parameters(),
+        lr=args.learning_rate,
+        betas=(0.9, 0.999),
+        weight_decay=1e-2,
+        eps=1e-8,
+    )
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+    )
+
+    unet, optimizer, dataloader, lr_scheduler = accelerator.prepare(
+        unet, 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)
+
+    global_step = 0
+    data_iterator = iter(dataloader)
+    while global_step < args.max_train_steps:
+        try:
+            batch = next(data_iterator)
+        except StopIteration:
+            data_iterator = iter(dataloader)
+            batch = next(data_iterator)
+
+        with accelerator.accumulate(unet):
+            pixel_values = batch["pixel_values"].to(
+                accelerator.device, dtype=weight_dtype, non_blocking=True
+            )
+            masked_pixel_values = batch["masked_pixel_values"].to(
+                accelerator.device, dtype=weight_dtype, non_blocking=True
+            )
+            mask = batch["mask"].to(
+                accelerator.device, dtype=weight_dtype, non_blocking=True
+            )
+            with torch.no_grad():
+                latents = vae.encode(pixel_values).latent_dist.sample()
+                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]
+
+            noise = torch.randn_like(latents)
+            timesteps = torch.randint(
+                0,
+                noise_scheduler.config.num_train_timesteps,
+                (latents.shape[0],),
+                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(
+                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)
+            else:
+                raise ValueError(
+                    f"Unsupported prediction type: {noise_scheduler.config.prediction_type}"
+                )
+            loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+            accelerator.backward(loss)
+            if accelerator.sync_gradients:
+                accelerator.clip_grad_norm_(unet.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:
+                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}"
+                )
+            accelerator.log(
+                {"train_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]},
+                step=global_step,
+            )
+
+    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"
+        )
+        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.end_training()
+
+
+if __name__ == "__main__":
+    main()