Spaces:
Running on Zero
Running on Zero
File size: 2,675 Bytes
e967c28 | 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 | """
Part 1 entry point: generate images for selected prompts with a chosen
model, one image at a time (no batching, no analysis during generation).
Usage:
from src.prompts import load_prompt_bank, select_prompts
from src.providers import PROVIDERS, MockProvider
from src.generate import run_generation
bank = load_prompt_bank("data/prompt_bank.csv")
prompts = select_prompts(bank, categories=["occupation"], n_per_category=5, seed=1)
provider = MockProvider() # or PROVIDERS["FX"]()
manifest_path, df = run_generation(
prompts=prompts,
provider=provider,
images_per_prompt=10,
output_dir="runs/my_first_audit",
)
"""
import os
from .manifest import ManifestWriter
def run_generation(prompts, provider, images_per_prompt, output_dir, on_progress=None):
"""
prompts: list[Prompt] from select_prompts()
provider: an instance of a ModelProvider subclass
images_per_prompt: int, 5-100
output_dir: where images + manifest chunks/final file are written
on_progress: optional callback(done, total) -- e.g. for a Gradio progress bar
Returns: (manifest_path, manifest_dataframe)
"""
if not (5 <= images_per_prompt <= 100):
raise ValueError("images_per_prompt must be between 5 and 100")
images_dir = os.path.join(output_dir, "images")
os.makedirs(images_dir, exist_ok=True)
writer = ManifestWriter(output_dir)
total = len(prompts) * images_per_prompt
done = 0
for prompt in prompts:
prompt_dir = os.path.join(images_dir, prompt.category, f"{prompt.prompt_id:03d}")
os.makedirs(prompt_dir, exist_ok=True)
for image_index in range(1, images_per_prompt + 1):
image_id = f"{provider.model_key}_{prompt.category}_{prompt.prompt_id:03d}_{image_index:03d}"
if writer.is_done(image_id):
done += 1
if on_progress:
on_progress(done, total)
continue
image_bytes = provider.generate(prompt.prompt_text)
file_path = os.path.join(prompt_dir, f"{image_index:03d}.png")
with open(file_path, "wb") as f:
f.write(image_bytes)
writer.add_row(
model=provider.model_key,
category=prompt.category,
prompt_id=prompt.prompt_id,
prompt_text=prompt.prompt_text,
image_index=image_index,
file_path=file_path,
)
done += 1
if on_progress:
on_progress(done, total)
manifest_path, df = writer.finalize()
return manifest_path, df
|