# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "torch",
#   "diffusers>=0.31",
#   "transformers>=4.44",
#   "peft>=0.13",
#   "huggingface_hub>=0.24",
#   "pillow",
# ]
# ///
"""Stage 3 of training: render the held-out eval prompts from a checkpoint.

    # baseline, before training - run this FIRST
    hf jobs uv run --flavor a10g-small --timeout 1h --secrets HF_TOKEN \\
        -e CHECKPOINT=base training/sample_job.py

    # any checkpoint afterwards
    hf jobs uv run --flavor a10g-small --timeout 1h --secrets HF_TOKEN \\
        -e CHECKPOINT=checkpoint-2000 training/sample_job.py

Validation loss on a diffusion fine-tune correlates weakly with how the
output actually looks, so it cannot be the only signal. This renders the 50
prompts in eval_prompts.json - drawn from validation groups, so the model has
never trained on them - at a fixed seed, and pushes a contact sheet to the
model repo.

Run CHECKPOINT=base before training starts. Without a baseline you cannot
tell improvement from drift.
"""

import json
import os
from pathlib import Path

import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import HfApi, hf_hub_download
from PIL import Image

DATASET_REPO = "whosouravsharma/text-to-image-diffusiondb-2M"
DATASET_REVISION = os.environ.get("DATASET_REVISION", "v2-clean")
MODEL_REPO = os.environ.get("MODEL_REPO", "whosouravsharma/diffusiondb-sd15-lora")

BASE_MODEL = "runwayml/stable-diffusion-v1-5"

# "base" renders vanilla SD 1.5 with no adapter - the comparison baseline.
CHECKPOINT = os.environ.get("CHECKPOINT", "base")
CHECKPOINT_DIR = "checkpoints"

STEPS = int(os.environ.get("STEPS", 30))
GUIDANCE = float(os.environ.get("GUIDANCE", 7.5))
SEED = int(os.environ.get("SEED", 42))
COLUMNS = 5
THUMB = 320


def contact_sheet(images: list[Image.Image]) -> Image.Image:
    rows = (len(images) + COLUMNS - 1) // COLUMNS
    sheet = Image.new("RGB", (COLUMNS * THUMB, rows * THUMB), (18, 18, 22))
    for index, image in enumerate(images):
        sheet.paste(
            image.resize((THUMB, THUMB), Image.LANCZOS),
            ((index % COLUMNS) * THUMB, (index // COLUMNS) * THUMB),
        )
    return sheet


def main() -> None:
    if not torch.cuda.is_available():
        raise SystemExit("No GPU. Run this on a GPU flavor, e.g. a10g-small.")

    prompts = json.loads(Path(hf_hub_download(
        DATASET_REPO, "eval_prompts.json", repo_type="dataset",
        revision=DATASET_REVISION,
    )).read_text())
    print(f"{len(prompts)} eval prompts from {DATASET_REPO}@{DATASET_REVISION}")

    pipeline = StableDiffusionPipeline.from_pretrained(
        BASE_MODEL, torch_dtype=torch.float16, safety_checker=None,
    ).to("cuda")
    pipeline.set_progress_bar_config(disable=True)

    if CHECKPOINT != "base":
        print(f"Loading LoRA {CHECKPOINT} ...")
        pipeline.load_lora_weights(
            MODEL_REPO, weight_name="pytorch_lora_weights.safetensors",
            subfolder=f"{CHECKPOINT_DIR}/{CHECKPOINT}",
        )

    images = []
    for index, prompt in enumerate(prompts):
        # Same seed per prompt index across checkpoints, so differences come
        # from the weights rather than from the noise.
        generator = torch.Generator("cuda").manual_seed(SEED + index)
        images.append(pipeline(
            prompt, num_inference_steps=STEPS,
            guidance_scale=GUIDANCE, generator=generator,
        ).images[0])
        if (index + 1) % 10 == 0:
            print(f"  {index + 1}/{len(prompts)}", flush=True)

    out = Path("samples") / CHECKPOINT
    out.mkdir(parents=True, exist_ok=True)
    for index, image in enumerate(images):
        image.save(out / f"{index:03}.png")
    contact_sheet(images).save(out / "grid.jpg", quality=92)
    (out / "prompts.json").write_text(json.dumps(
        {"checkpoint": CHECKPOINT, "steps": STEPS, "guidance": GUIDANCE,
         "seed": SEED, "prompts": prompts}, indent=2))

    # MODEL_REPO is created ahead of time, with its card, rather than by the job.
    api = HfApi()
    api.upload_folder(
        folder_path=str(out), path_in_repo=f"samples/{CHECKPOINT}",
        repo_id=MODEL_REPO, repo_type="model",
        commit_message=f"Eval samples for {CHECKPOINT}",
    )
    print(f"\nPushed to https://huggingface.co/{MODEL_REPO}"
          f"/tree/main/samples/{CHECKPOINT}")


if __name__ == "__main__":
    main()