File size: 4,401 Bytes
1ab3690 d6f5ba0 1ab3690 d6f5ba0 1ab3690 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | # /// 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()
|