dataset.py

"""Datasets and prompts used by the custom StyleDrop trainer."""

import json
from pathlib import Path

from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms


GRID_SUBJECTS = (
    "A flying pigeon ",
    "A bicycle ",
    "A tree in the city ",
    "A running horse ",
    "A fruit bowl ",
    "A brown donkey ",
    "A castle on the beach ",
    "A human sitting on a bench ",
    "A fish in an aquarium ",
    "A mountain scene ",
)

EVALUATION_SUBJECTS = (
    "A flying pigeon ",
    "A bicycle ",
    "A tree in the city ",
    "A running horse ",
    "A fruit bowl ",
    "A brown donkey ",
    "A castle on the beach ",
    "A human sitting on a bench ",
    "A fish in an aquarium ",
    "A mountain scene ",
    "A butterfly ",
    "A house ",
    "A car on the street ",
    "A deer in the field ",
    "A banana on the table ",
    "A black cat ",
    "A church on the street ",
    "A human walking in the forest ",
    "A dolphin in water ",
    "A city skyline ",
    "A dog ",
    "A moose ",
)


def build_style_prompts(style, subjects):
    return [f"{subject} {style}" for subject in subjects]


def clamp_images(images):
    """Convert decoded images to the range expected by torchvision."""
    return images.clamp_(0.0, 1.0)


class StyleTrainingDataset(Dataset):
    """Read reference images and prompts from the StyleDrop JSON format."""

    def __init__(self, metadata_file, samples_per_epoch=24):
        metadata_path = Path(metadata_file)
        with metadata_path.open(encoding="utf-8") as file:
            metadata = json.load(file)

        if not metadata:
            raise ValueError(f"No reference images found in {metadata_path}")

        self.image_paths = []
        self.prompts = []
        styles = set()

        for filename, description in metadata.items():
            if not isinstance(description, list) or len(description) != 2:
                raise ValueError(
                    f"{filename!r} must map to [subject, style]"
                )
            subject, style = description
            image_path = metadata_path.parent / filename
            if not image_path.is_file():
                raise FileNotFoundError(
                    f"Reference image does not exist: {image_path}"
                )

            self.image_paths.append(image_path)
            self.prompts.append(f"{subject} {style}")
            styles.add(style)

        if len(styles) != 1:
            raise ValueError(
                "A training run must contain exactly one shared style; "
                f"found {sorted(styles)}"
            )

        self.style = styles.pop()
        self.samples_per_epoch = samples_per_epoch
        self.transform = transforms.Compose(
            [
                transforms.Resize((256, 256)),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
            ]
        )

    @property
    def num_reference_images(self):
        return len(self.image_paths)

    def __getitem__(self, index):
        reference_index = index % self.num_reference_images
        image = Image.open(
            self.image_paths[reference_index]
        ).convert("RGB")
        return (
            self.transform(image),
            self.prompts[reference_index],
        )

    def __len__(self):
        return self.samples_per_epoch