pruna-client 0.0.7

Official Pruna API client for synchronous and asynchronous image generation and editing

Pruna Client

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

Installation

uv add pruna-client

Quick Start

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.

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

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

Text to Image

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

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

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)

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

File Upload

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

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.

response = client.poll_status(response=response)

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

Closing the Client

client.close()

Running Tests

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.