scripts/dataset/generate_auraflow_synthetic.py

"""
Generate synthetic images with AuraFlow v0.3.

License
-------
`fal/AuraFlow-v0.3` is released under Apache-2.0. Stage 3A approved it as an
independent rectified-flow training source. See `NOTICES.md` before running a
real generation job.

Usage
-----
    python scripts/dataset/generate_auraflow_synthetic.py \\
        --out data/raw/ai_generated/auraflow-v0.3 \\
        --count 20000 \\
        --prompts scripts/dataset/prompts.txt

Dry run, no model import/download:

    python scripts/dataset/generate_auraflow_synthetic.py \\
        --out data/raw/ai_generated/auraflow-v0.3 \\
        --count 3 \\
        --dry-run

Hardware
--------
AuraFlow 1024x1024 generation should be run on a CUDA GPU with enough VRAM for
half-precision inference. Start with a small smoke run on the target GPU before
renting a long job.

Idempotency
-----------
Each generated file is named by a hash of (prompt, seed), so re-runs skip
already-generated images. Crash-resume works automatically.
"""
from __future__ import annotations

import argparse
import csv
import random
from pathlib import Path

from generation_utils import (
    APPROVED_GENERATORS,
    STAGE3A_MANIFEST_FIELDS,
    choose_prompt,
    image_dimensions,
    infer_data_root,
    load_prompts,
    manifest_row,
    next_seed,
    sha256_file,
    stable_image_key,
)


FALLBACK_PROMPTS = [
    "a realistic phone photo of a ceramic mug beside a laptop",
    "a casual snapshot of a parking lot after a summer storm",
    "a natural light photo of houseplants on a crowded windowsill",
    "a slightly blurry photo of a folded jacket on a cafe chair",
]


AURAFLOW_SPEC = APPROVED_GENERATORS["auraflow-v0.3"]


def _manifest_path(out_dir: Path) -> Path:
    return out_dir.parent / "auraflow_manifest.csv"


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--out", type=Path, required=True, help="Output directory")
    parser.add_argument(
        "--data-root",
        type=Path,
        default=None,
        help=(
            "Dataset root for manifest paths. Defaults to the parent of the "
            "'raw' path segment in --out."
        ),
    )
    parser.add_argument("--count", type=int, default=20_000)
    parser.add_argument(
        "--prompts",
        type=Path,
        default=None,
        help="Optional file with one prompt per line",
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--steps", type=int, default=30)
    parser.add_argument("--guidance-scale", type=float, default=3.5)
    parser.add_argument("--width", type=int, default=1024)
    parser.add_argument("--height", type=int, default=1024)
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate prompt/seed/output planning without loading AuraFlow or writing files",
    )
    args = parser.parse_args()

    out_dir = args.out.resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    data_root = (
        args.data_root.resolve()
        if args.data_root is not None
        else infer_data_root(out_dir)
    )
    manifest_path = _manifest_path(out_dir)

    prompts = load_prompts(args.prompts, fallback_prompts=FALLBACK_PROMPTS)
    if args.prompts is not None:
        print(f"Loaded {len(prompts)} prompts from {args.prompts}")
    else:
        print(
            f"WARNING: --prompts not given; using {len(prompts)} built-in "
            "fallback prompts (smoke-test only, not enough diversity for "
            "a real training run)"
        )

    rng = random.Random(args.seed)

    if args.dry_run:
        print(
            "Dry run: AuraFlow pipeline will not be loaded and no images "
            "will be written."
        )
        for i in range(args.count):
            prompt = choose_prompt(rng, prompts)
            seed = next_seed(rng)
            key = stable_image_key(prompt, seed)
            dst = out_dir / f"{key}.png"
            print(f"  {i + 1:04d}: seed={seed} path={dst} prompt={prompt!r}")
        print(f"Dry run complete. Planned manifest fragment: {manifest_path}")
        return

    print("Loading AuraFlow pipeline (large download on first run)...")
    import torch
    from diffusers import AuraFlowPipeline

    pipe = AuraFlowPipeline.from_pretrained(
        AURAFLOW_SPEC.model_id,
        torch_dtype=torch.float16,
    )
    pipe.to("cuda")
    # AuraFlow VAE dtype-mismatch fix.
    # Why: AuraFlow's VAE has biases that don't survive `torch_dtype=torch.float16`
    # cleanly. Diffusers' internal `upcast_vae()` path (now deprecated for AuraFlow)
    # only partially upcasts, leaving conv biases stranded in fp32 while inputs are
    # fp16 -> RuntimeError "Input type (c10::Half) and bias type (float) should be
    # the same" during `vae.decode`.
    # Fix: cast the whole VAE to fp32, AND monkey-patch `vae.decode` to cast the
    # latent input to match. fp32 VAE adds ~500 MB memory and ~10-20% time to the
    # decode step (negligible on 48 GB cards).
    pipe.vae = pipe.vae.to(dtype=torch.float32)
    _orig_vae_decode = pipe.vae.decode
    def _decode_with_dtype_cast(z, *args, **kwargs):
        z = z.to(pipe.vae.dtype)
        return _orig_vae_decode(z, *args, **kwargs)
    pipe.vae.decode = _decode_with_dtype_cast

    rows: list[dict] = []
    for i in range(args.count):
        prompt = choose_prompt(rng, prompts)
        seed = next_seed(rng)
        key = stable_image_key(prompt, seed)
        dst = out_dir / f"{key}.png"

        if not dst.exists():
            generator = torch.Generator("cuda").manual_seed(seed)
            image = pipe(
                prompt=prompt,
                num_inference_steps=args.steps,
                guidance_scale=args.guidance_scale,
                width=args.width,
                height=args.height,
                generator=generator,
            ).images[0]
            image.save(dst, format="PNG")

        width, height = image_dimensions(dst)
        rows.append(
            manifest_row(
                path=dst.relative_to(data_root).as_posix(),
                cls="ai_generated",
                spec=AURAFLOW_SPEC,
                sha256=sha256_file(dst),
                prompt=prompt,
                seed=seed,
                width=width,
                height=height,
                generation_params={
                    "steps": args.steps,
                    "guidance_scale": args.guidance_scale,
                    "width": args.width,
                    "height": args.height,
                    "pipeline": AURAFLOW_SPEC.pipeline,
                },
            )
        )

        if (i + 1) % 100 == 0:
            print(f"  generated {i + 1}/{args.count}")

    with manifest_path.open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=STAGE3A_MANIFEST_FIELDS)
        writer.writeheader()
        writer.writerows(rows)

    print(f"Done. {len(rows)} images. Manifest fragment: {manifest_path}")


if __name__ == "__main__":
    main()