Spaces:
Running on Zero
Running on Zero
| """ | |
| 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 | |