Metadata-Version: 2.4
Name: pruna-client
Version: 0.0.8.1
Summary: Official Pruna API client for synchronous and asynchronous image generation and editing
Requires-Python: >=3.11
Requires-Dist: cachetools>=5.3.0
Requires-Dist: httpx[http2]>=0.27.0
Requires-Dist: pillow>=10.0.0
Requires-Dist: rich>=13.0.0
Provides-Extra: dev
Requires-Dist: pytest-mock>=3.12.0; extra == 'dev'
Requires-Dist: pytest>=8.0.0; extra == 'dev'
Description-Content-Type: text/markdown

# Pruna Client

Official Pruna API client for image and video generation and editing.

## Installation

```bash
uv add pruna-client
```

## Quick Start

```python
from pruna_client import PrunaClient
from pruna_client.models import PredictionStatus

# Initialize client (uses PRUNA_API_KEY env var if api_key not provided)
client = PrunaClient(api_key="your_api_key")

# Generate image
response = client.generate_text_to_image(
    model="p-image",
    prompt="A beautiful sunset over a calm ocean",
    sync=True,
)

# Access the generated content via generation_url
if response.status == PredictionStatus.SUCCEEDED:
    generation_url = response.response.get("generation_url")
    if generation_url:
        # Download the generated image
        image_bytes = client.download_content(generation_url)
        with open("generated_image.jpg", "wb") as f:
            f.write(image_bytes)
```

## Accessing Generated Content

**Important:** All successful generation responses contain a `generation_url` in `response.response["generation_url"]` that points to the generated content (image or video). You can download this content using the `download_content()` method.

```python
from pruna_client.models import PredictionStatus

# After generating content
if response.status == PredictionStatus.SUCCEEDED:
    generation_url = response.response.get("generation_url")
    if generation_url:
        # Download the generated content
        content = client.download_content(generation_url)
        with open("output.jpg", "wb") as f:  # or .mp4 for videos
            f.write(content)
```

## Basic Usage

### General Generation

We support both batch and single generation for any model. You only need to call the `generate` or `generate_batch` method with the required parameters as shown on [Pruna API Reference](https://ac25aba50212.eu.kongportals.com/)

```python
from pruna_client.models import PredictionStatus

response = client.generate(
    model="p-image",
    input={"prompt": "A beautiful sunset over a calm ocean"},
    sync=True
)

# Access the generated content via generation_url
if response.status == PredictionStatus.SUCCEEDED:
    generation_url = response.response.get("generation_url")
    if generation_url:
        content = client.download_content(generation_url)
        with open("output.jpg", "wb") as f:
            f.write(content)

# or batch
responses = client.generate_batch(
    requests=[
        {"model": "p-image", "input": {"prompt": "A sunset"}, "sync": True},
        {"model": "p-image", "input": {"prompt": "A sunrise"}, "sync": True},
    ]
)

# Access generation URLs from batch responses
for i, response in enumerate(responses):
    if response.status == PredictionStatus.SUCCEEDED:
        generation_url = response.response.get("generation_url")
        if generation_url:
            content = client.download_content(generation_url)
            with open(f"output_{i}.jpg", "wb") as f:
                f.write(content)
```

### Specific Models

We support specific models for image and video generation and editing. You only need to call the `generate_text_to_image`, `generate_image_edit`, `generate_text_to_video`, `generate_video_edit`, `generate_image_to_video` method with the required parameters as shown on [Pruna API Reference](https://ac25aba50212.eu.kongportals.com/)

### Text to Image

```python
from pruna_client.models import PredictionStatus

response = client.generate_text_to_image(
    model="p-image",
    prompt="A beautiful sunset over a calm ocean",
    sync=True
)

# The response contains a generation_url with the generated content
if response.status == PredictionStatus.SUCCEEDED:
    generation_url = response.response.get("generation_url")
    if generation_url:
        # Download the generated image
        image_bytes = client.download_content(generation_url)
        with open("generated_image.jpg", "wb") as f:
            f.write(image_bytes)
```

### Image Editing

```python
from pruna_client.models import PredictionStatus

response = client.generate_image_edit(
    model="p-image-edit",
    prompt="Edit the image to make it more beautiful",
    images=["path/to/image.png"],
    sync=True
)

# Access the generated content via generation_url
if response.status == PredictionStatus.SUCCEEDED:
    generation_url = response.response.get("generation_url")
    if generation_url:
        image_bytes = client.download_content(generation_url)
        with open("edited_image.jpg", "wb") as f:
            f.write(image_bytes)
```

### Text to Video

```python
from pruna_client.models import PredictionStatus

response = client.generate_text_to_video(
    model="wan-t2v",
    prompt="A beautiful sunset over a calm ocean",
    sync=False
)

# Poll for completion and access generation_url
final_response = client.poll_status(response=response)
if final_response.status == PredictionStatus.SUCCEEDED:
    generation_url = final_response.response.get("generation_url")
    if generation_url:
        video_bytes = client.download_content(generation_url)
        with open("generated_video.mp4", "wb") as f:
            f.write(video_bytes)
```

## Async Usage

The client supports async operations for better performance when making multiple requests or integrating with async applications.

### General Generation (Async)

```python
import asyncio
from pruna_client import PrunaClient
from PIL import Image

async def main():
    from pruna_client.models import PredictionStatus
    
    client = PrunaClient(api_key="your_api_key")
    
    response = await client.agenerate(
        model="p-image",
        input={"prompt": "A beautiful sunset over a calm ocean"},
        sync=True
    )
    
    # Access the generated content via generation_url
    if response.status == PredictionStatus.SUCCEEDED:
        generation_url = response.response.get("generation_url")
        if generation_url:
            content = client.download_content(generation_url)
            with open("output.jpg", "wb") as f:
                f.write(content)
    
    await client.aclose()

asyncio.run(main())

# or batch
async def batch_example():
    from pruna_client.models import PredictionStatus
    
    client = PrunaClient(api_key="your_api_key")
    responses = await client.agenerate_batch(
        requests=[
            {"model": "p-image", "input": {"prompt": "A beautiful sunset over a calm ocean"}, "sync": True},
            {"model": "p-image", "input": {"prompt": "A beautiful sunrise over a calm ocean"}, "sync": True},
        ]
    )
    
    # Access generation URLs from batch responses
    for i, response in enumerate(responses):
        if response.status == PredictionStatus.SUCCEEDED:
            generation_url = response.response.get("generation_url")
            if generation_url:
                content = client.download_content(generation_url)
                with open(f"output_{i}.jpg", "wb") as f:
                    f.write(content)
    
    await client.aclose()

asyncio.run(batch_example())
```

## Additional methods

We support additional methods for file upload, polling status, and closing the client. You only need to call the `upload_file`, `poll_status`, `close` method with the required parameters as shown on [Pruna API Reference](https://ac25aba50212.eu.kongportals.com/)

### File Upload

We support `string`, `pathlib.Path`, `PIL.Image.Image`, and `bytes` as input.

```python
url = client.upload_file("path/to/image.png")

# or batch
urls = client.upload_file_batch(["path/to/image.png", "path/to/image2.png"])
```

### Polling Status

This method is used to poll the status of a generation request and return the final response.

```python
response = client.poll_status(response=response)

# or status URL
response = client.poll_status(status_url=response.response["get_url"])
```

### Closing the Client

```python
client.close()
```

## Running Tests

```bash
uv run pytest tests/integration/test_general.py -v
uv run pytest tests/integration/test_text_to_image.py -v
uv run pytest tests/integration/test_image_edit.py -v
uv run pytest tests/integration/test_text_to_video.py -v
uv run pytest tests/integration/test_video_edit.py -v
uv run pytest tests/integration/test_image_to_video.py -v
uv run pytest tests/integration/test_batch_generation.py -v
```

Tests require `PRUNA_API_KEY` environment variable to be set.
