prompt-tuners 0.1.0
Automatically tune prompts for image generation and image editing and text generation using DSPy
License: MIT
Requires Python: >=3.11
Keywords: ai, dspy, image-generation, llm, optimization, prompt-engineering, pydantic

# prompt-tuners

Iterative prompt refinement for image generation and editing using DSPy.

## Overview

`prompt-tuners` automatically refines prompts through iterative feedback loops, improving image generation and editing results by analyzing outputs and adjusting prompts until the desired outcome is achieved. It also learns prompt rewriting patterns from examples and can automatically train task-specific rewriters during iterations.

## Architecture

The library is built with a modular architecture:

- **`BasePromptTuner`** - Base class with shared logic for prompt refinement, template handling, learning, and evaluation
- **`TextToImageTuner`** - Task-specific tuner for text-to-image generation
- **`ImageEditTuner`** - Task-specific tuner for image editing
- **`ImageProvider`** - Abstract base class for image generation/editing providers
- **`ReplicateProvider`** - Default Replicate API provider implementation

## Quickstart

### Installation

```bash
uv pip install -e .
```

### Setup

Set your API keys in a `.env` file:

```bash
REPLICATE_API_TOKEN=your_token_here
OPENAI_API_KEY=your_key_here  # For DSPy LLM
```

### Basic Usage

```python
import dspy
from dotenv import load_dotenv
from prompt_tuners import TextToImageTuner, ImageEditTuner

load_dotenv()

# Configure DSPy
lm = dspy.LM(model="gpt-4o-mini", temperature=0.5)
dspy.configure(lm=lm)

# For text-to-image generation
text_tuner = TextToImageTuner(max_iterations=5)
iterations = text_tuner.generate_image(
    input_prompt="A serene autumn scene with fog"
)
image, final_prompt = iterations[-1]
print(f"Image URL: {image.url}")

# For image editing
edit_tuner = ImageEditTuner(max_iterations=5)
iterations = edit_tuner.edit_image(
    input_prompt="Make it more dramatic",
    image_inputs="path/to/image.png"
)
edited_image, edited_prompt = iterations[-1]
```

## Key Features

### Auto-Tuning

By default, `auto_tune=True` automatically collects (initial, final) prompt pairs during iterations and trains task-specific rewriters:

- **`text_to_image_rewriter`** - Trained from text-to-image examples
- **`image_edit_rewriter`** - Trained from image-edit examples
- **`text_generation_rewriter`** - General prompt rewriter (trained via `learn_text_generation()`)

Rewriters are automatically trained when at least 3 examples are collected.

### Learning from Examples

Train prompt rewriters from curated examples:

```python
# Learn general text generation patterns
general_examples = [
    ("a cat", "a photorealistic orange tabby cat, natural lighting, detailed fur"),
    ("a dog", "a golden retriever puppy in a sunny park, high detail"),
]
tuner.learn_text_generation(general_examples)

# Learn text-to-image patterns
text_examples = [
    ("a cat", "a photorealistic orange tabby cat, natural lighting, detailed fur"),
    ("a dog", "a golden retriever puppy in a sunny park, high detail"),
]
tuner.learn_text_to_image(text_examples)

# Learn image-edit patterns
edit_examples = [
    ("make it brighter", "enhance brightness, add warm lighting"),
    ("change background", "replace with serene forest, maintain subject"),
]
tuner.learn_image_edit(edit_examples)
```

### Evaluation Criteria

Customize evaluation criteria to guide prompt refinement. Criteria can be specified as a string or dictionary when initializing the tuner:

```python
# Using default criteria (prompt_adherence, general_aesthetics)
tuner = TextToImageTuner(max_iterations=5)

# Using custom string criteria
tuner = TextToImageTuner(
    max_iterations=5,
    evaluation_criteria="The image should be vibrant and colorful with high contrast"
)

# Using custom dictionary criteria
tuner = TextToImageTuner(
    max_iterations=5,
    evaluation_criteria={
        "color_vibrancy": "The image should use vibrant, saturated colors",
        "composition": "The image should follow the rule of thirds",
        "detail_level": "The image should show fine details and textures"
    }
)
```

The tuner evaluates images against these criteria during each iteration, providing scores (1-5) for each criterion and using them to guide prompt refinement.

### Custom Providers

The library uses a provider abstraction that allows you to use different image generation/editing backends. By default, it uses `ReplicateProvider`, but you can create custom providers or use alternative implementations.

#### Using the Default Replicate Provider

```python
from prompt_tuners import TextToImageTuner, ReplicateProvider

# Explicitly create a Replicate provider
provider = ReplicateProvider(token="your-token")
tuner = TextToImageTuner(provider=provider)
```

#### Creating a Custom Provider

You can create custom providers by implementing the `ImageProvider` interface:

```python
from prompt_tuners import ImageProvider, TextToImageTuner
import dspy

class MyCustomProvider(ImageProvider):
    """Custom provider for your image generation service."""
    
    def generate_image(self, prompt: str, **kwargs) -> dspy.Image:
        # Your custom image generation logic
        # Return a dspy.Image object
        image_url = your_api_call(prompt, **kwargs)
        return dspy.Image.from_url(image_url)
    
    def edit_image(
        self,
        prompt: str,
        image_inputs: Union[str, Path, List[Union[str, Path]]],
        **kwargs,
    ) -> dspy.Image:
        # Your custom image editing logic
        # Return a dspy.Image object
        image_url = your_edit_api_call(prompt, image_inputs, **kwargs)
        return dspy.Image.from_url(image_url)

# Use your custom provider
provider = MyCustomProvider()
tuner = TextToImageTuner(provider=provider, max_iterations=5)
results = tuner.generate_image("a beautiful landscape")
```

#### Provider-Specific Parameters

You can pass provider-specific parameters through `**kwargs`:

```python
# Replicate provider accepts model and other Replicate-specific params
results = tuner.generate_image(
    "a cat",
    model="stability-ai/sdxl",  # Replicate model
    num_inference_steps=50,      # Replicate-specific parameter
    guidance_scale=7.5            # Replicate-specific parameter
)
```

### Batch Processing

Process multiple prompts with automatic rewriter training:

```python
# Batch generation
results = tuner.generate_image_batch(
    input_prompts=["prompt 1", "prompt 2", "prompt 3"],
    auto_tune=True,  # Default
)

# Batch editing
results = tuner.edit_image_batch(
    input_prompts=["edit 1", "edit 2"],
    image_inputs="path/to/image.png",
    auto_tune=True,  # Default
)
```

### Saving and Loading

Save and load trained rewriters using DSPy's native methods:

```python
# Save rewriters
tuner.save_text_to_image_rewriter("text_rewriter.json", save_program=False)
tuner.save_image_edit_rewriter("edit_rewriter.json", save_program=False)
tuner.save_text_generation("text_generation.json", save_program=False)

# Load rewriters
tuner.load_text_to_image_rewriter("text_rewriter.json", save_program=False)
tuner.load_image_edit_rewriter("edit_rewriter.json", save_program=False)
tuner.load_text_generation("text_generation.json", save_program=False)
```

## API Reference

### Initialization

#### TextToImageTuner

- **`TextToImageTuner(max_iterations=5, evaluation_criteria=None, provider=None)`** - Initialize text-to-image tuner
  - `max_iterations`: Maximum number of refinement iterations (default: 5)
  - `evaluation_criteria`: Custom evaluation criteria. Can be:
    - `None`: Uses default criteria (prompt_adherence, general_aesthetics)
    - `str`: A string describing custom criteria
    - `dict`: Dictionary mapping criterion names to descriptions
  - `provider`: Optional `ImageProvider` instance (defaults to `ReplicateProvider`)

#### ImageEditTuner

- **`ImageEditTuner(max_iterations=5, evaluation_criteria=None, provider=None)`** - Initialize image edit tuner
  - `max_iterations`: Maximum number of refinement iterations (default: 5)
  - `evaluation_criteria`: Custom evaluation criteria. Can be:
    - `None`: Uses default criteria (prompt_adherence, general_aesthetics)
    - `str`: A string describing custom criteria
    - `dict`: Dictionary mapping criterion names to descriptions
  - `provider`: Optional `ImageProvider` instance (defaults to `ReplicateProvider`)

#### ReplicateProvider

- **`ReplicateProvider(token=None)`** - Initialize Replicate API provider
  - `token`: Replicate API token (optional, reads from `REPLICATE_API_TOKEN` env var if not provided)

### Core Methods

#### TextToImageTuner

- **`generate_image(input_prompt, model, template_variables, auto_tune=True, **kwargs)`** - Generate images with iterative refinement
  - `input_prompt`: Initial prompt for image generation
  - `model`: Model identifier (passed to provider, default: "prunaai/p-image")
  - `template_variables`: Optional dict for template placeholders
  - `auto_tune`: If True, collects examples for training
  - `**kwargs`: Additional provider-specific arguments
- **`generate_image_batch(input_prompts, model, template_variables, auto_tune=True, **kwargs)`** - Batch image generation

#### ImageEditTuner

- **`edit_image(input_prompt, image_inputs, model, template_variables, auto_tune=True, **kwargs)`** - Edit images with iterative refinement
  - `input_prompt`: Initial prompt for image editing
  - `image_inputs`: Image(s) to edit (single or list)
  - `model`: Model identifier (passed to provider, default: "prunaai/p-image-edit")
  - `template_variables`: Optional dict for template placeholders
  - `auto_tune`: If True, collects examples for training
  - `**kwargs`: Additional provider-specific arguments
- **`edit_image_batch(input_prompts, image_inputs, model, template_variables, auto_tune=True, **kwargs)`** - Batch image editing

#### BasePromptTuner (All Tuners)

- **`optimize_prompt(initial_prompt)`** - Optimize a prompt using the learned `text_generation_rewriter`

### Learning Methods

- **`learn_text_generation(examples)`** - Learn general text generation patterns (trains `text_generation_rewriter`)
- **`learn_text_to_image(examples)`** - Learn text-to-image specific patterns (trains `text_to_image_rewriter`)
- **`learn_image_edit(examples)`** - Learn image-edit specific patterns (trains `image_edit_rewriter`)

### Save/Load Methods

- **`save_text_generation(filepath, save_program=False)`** - Save text generation rewriter
- **`load_text_generation(filepath, save_program=False, allow_pickle=False)`** - Load text generation rewriter
- **`save_text_to_image_rewriter(filepath, save_program=False)`** - Save text-to-image rewriter
- **`load_text_to_image_rewriter(filepath, save_program=False, allow_pickle=False)`** - Load text-to-image rewriter
- **`save_image_edit_rewriter(filepath, save_program=False)`** - Save image-edit rewriter
- **`load_image_edit_rewriter(filepath, save_program=False, allow_pickle=False)`** - Load image-edit rewriter

## Examples

See the `examples/` directory for complete examples:

- **`text_to_image_without_template.py`** - Basic image generation
- **`text_to_image_with_template.py`** - Image generation with template placeholders
- **`image_edit_without_template.py`** - Basic image editing
- **`image_edit_with_template.py`** - Image editing with template placeholders
- **`learn_text_to_image_pattern.py`** - Learning text-to-image patterns from examples
- **`learn_image_edit_pattern.py`** - Learning image-edit patterns from examples
- **`auto_tune_example.py`** - Auto-tuning during iterations
- **`batch_with_auto_tune.py`** - Batch processing with auto-tuning