populate_test_data.py

import random
import time
import argparse
from PIL import Image
from prompt_history import HistoryManager
from generation_metadata import GenerationMetadata

def populate_test_data(num_entries=50, root=None):
    hm = HistoryManager(storage_root=root)
    for i in range(num_entries):
        # Create a dummy image
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
        img = Image.new('RGB', (100, 100), color=color)

        prompt = f"Test prompt {i}: A beautiful landscape with mountains and rivers"
        negative_prompt = f"Negative {i}: ugly, blurry, low quality"
        seed = random.randint(0, 1000000)
        width = random.choice([512, 768, 1024])
        height = random.choice([512, 768, 1024])
        guidance_scale = round(random.uniform(5.0, 10.0), 1)
        guidance_rescale = round(random.uniform(0.0, 1.0), 1)
        num_inference_steps = random.randint(10, 50)
        model_repo_id = "damian0815/loam_dfn5b"
        device = "cpu"

        input_metadata = GenerationMetadata(
            prompt=prompt,
            negative_prompt=negative_prompt,
            seed=seed,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            guidance_rescale=guidance_rescale,
            num_inference_steps=num_inference_steps,
            model_repo_id=model_repo_id,
            device=device
        )

        hm.save_generation(img, input_metadata)

        print(f"Added test entry {i+1}/{num_entries}")
        time.sleep(1)  # Ensure different timestamps

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Populate test data for prompt history.")
    parser.add_argument('--root', required=True, help='Root directory for storage (e.g., /data/loam_history)')
    parser.add_argument('--num_entries', type=int, default=50, help='Number of test entries to generate (default: 50)')
    args = parser.parse_args()
    populate_test_data(num_entries=args.num_entries, root=args.root)