scripts/dataset/generation_utils.py

"""
Shared helpers for Stage 3A dataset generation scripts.

Keep this module dependency-light: it is used on GPU dataset boxes, not in the
inference container. Generator-specific scripts should own their pipeline setup;
this module only centralizes the boring parts that must stay consistent across
sources.
"""
from __future__ import annotations

import hashlib
import json
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Sequence


@dataclass(frozen=True)
class GeneratorSpec:
    """License and metadata constants for one approved generator source."""

    source: str
    generator: str
    model_family: str
    model_id: str
    license: str
    license_url: str
    pipeline: str


APPROVED_GENERATORS: dict[str, GeneratorSpec] = {
    "flux.1-schnell": GeneratorSpec(
        source="flux.1-schnell",
        generator="flux.1-schnell",
        model_family="rectified_flow",
        model_id="black-forest-labs/FLUX.1-schnell",
        license="Apache-2.0",
        license_url="https://www.apache.org/licenses/LICENSE-2.0",
        pipeline="FluxPipeline",
    ),
    "sdxl": GeneratorSpec(
        source="sdxl",
        generator="sdxl",
        model_family="diffusion_unet",
        model_id="stabilityai/stable-diffusion-xl-base-1.0",
        license="CreativeML Open RAIL++-M",
        license_url=(
            "https://huggingface.co/stabilityai/"
            "stable-diffusion-xl-base-1.0/blob/main/LICENSE.md"
        ),
        pipeline="StableDiffusionXLPipeline",
    ),
    "sd3.5-medium": GeneratorSpec(
        source="sd3.5-medium",
        generator="sd3.5-medium",
        model_family="diffusion_transformer",
        model_id="stabilityai/stable-diffusion-3.5-medium",
        license="Stability AI Community License",
        license_url="https://stability.ai/license",
        pipeline="StableDiffusion3Pipeline",
    ),
    "auraflow-v0.3": GeneratorSpec(
        source="auraflow-v0.3",
        generator="auraflow-v0.3",
        model_family="rectified_flow",
        model_id="fal/AuraFlow-v0.3",
        license="Apache-2.0",
        license_url="https://www.apache.org/licenses/LICENSE-2.0",
        pipeline="AuraFlowPipeline",
    ),
}

STAGE3A_MANIFEST_FIELDS = [
    "path",
    "class",
    "source",
    "license",
    "license_url",
    "sha256",
    "generator",
    "model_family",
    "model_id",
    "prompt",
    "seed",
    "width",
    "height",
    "generation_params_json",
]


def load_prompts(
    path: Path | None,
    *,
    fallback_prompts: Sequence[str] | None = None,
) -> list[str]:
    """Read non-blank, non-comment prompts or return a fallback prompt list."""
    if path is None:
        if fallback_prompts is None:
            raise ValueError("No prompts path provided and no fallback prompts given")
        prompts = [p.strip() for p in fallback_prompts if p.strip()]
        if not prompts:
            raise ValueError("Fallback prompts list is empty")
        return prompts

    prompts: list[str] = []
    with path.open(encoding="utf-8") as fh:
        for line in fh:
            stripped = line.strip()
            if not stripped or stripped.startswith("#"):
                continue
            prompts.append(stripped)
    if not prompts:
        raise ValueError(f"No prompts found in {path}")
    return prompts


def choose_prompt(rng: random.Random, prompts: Sequence[str]) -> str:
    """Pick one prompt with the caller's deterministic RNG."""
    if not prompts:
        raise ValueError("prompts must not be empty")
    return rng.choice(list(prompts))


def next_seed(rng: random.Random) -> int:
    """Return a deterministic positive seed compatible with torch generators."""
    return rng.randint(0, 2**31 - 1)


def stable_image_key(prompt: str, seed: int) -> str:
    """Return the legacy content-addressed key for generated image files."""
    return hashlib.sha256(f"{prompt}|{seed}".encode()).hexdigest()[:24]


def infer_data_root(out_dir: Path) -> Path:
    """Infer the dataset root from an output path containing a `raw` segment."""
    resolved = out_dir.resolve()
    parts = resolved.parts
    if "raw" not in parts:
        raise ValueError(
            f"Cannot infer data root from {out_dir}; pass --data-root explicitly"
        )
    raw_index = parts.index("raw")
    if raw_index == 0:
        raise ValueError(
            f"Cannot infer data root from {out_dir}; pass --data-root explicitly"
        )
    return Path(*parts[:raw_index])


def sha256_file(path: Path) -> str:
    """Return the SHA-256 digest for a file."""
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def image_dimensions(path: Path) -> tuple[int, int]:
    """Return image width and height without keeping the image open."""
    from PIL import Image

    with Image.open(path) as image:
        return image.size


def params_json(params: dict[str, Any]) -> str:
    """Serialize generation parameters consistently for manifest rows."""
    return json.dumps(params, sort_keys=True, separators=(",", ":"))


def manifest_row(
    *,
    path: str,
    cls: str,
    spec: GeneratorSpec,
    sha256: str,
    prompt: str = "",
    seed: int | str = "",
    width: int | str = "",
    height: int | str = "",
    generation_params: dict[str, Any] | None = None,
    extra: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Build a Stage 3A-compatible manifest row for one generated image."""
    row: dict[str, Any] = {
        "path": path,
        "class": cls,
        "source": spec.source,
        "license": spec.license,
        "license_url": spec.license_url,
        "sha256": sha256,
        "generator": spec.generator,
        "model_family": spec.model_family,
        "model_id": spec.model_id,
        "prompt": prompt,
        "seed": seed,
        "width": width,
        "height": height,
        "generation_params_json": params_json(generation_params or {}),
    }
    if extra:
        row.update(extra)
    return row