""" 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()