Image Generation and Editing

Parasail currently supports batch image editing and generation using the OmniGen Model from VectorSpace Lab:

https://huggingface.co/Shitao/OmniGen-v1

This model can be accessed using both the OpenAI compatible interface API Reference and our Batch helper library Quick startthat wraps around the OpenAI client. The set of prompts are sent via JSONL files where each line is a prompt. The image is encoded using Base64 format (not Data URLS like "data:", rather the raw Base64 image data). The batch processing results are also sent via JSONL with Baes64 encoded output images. The batch submission files have a 500MB limit and

The following body parameters are currently supported:

Parameter	Description
model	Currently only "Shitao/OmniGen-v1" has been verified to work
prompt	The prompt to guide the generation. The prompts should reference the images as follows: "A man in a black shirt is reading a book. The man is the right man in <img><|image_1|></img>"
size	Formatted as "WxH" in pixels, e.g . "1024x1024"
image	A list of Base64-encoded JPGs or PNGs.
response_format	currently only "b64_json" is supported

Here is example code that uses our batch helper library to run through images in an input directory and save them to an output directory. The batch library has many additional useful features to build processing pipelines which can be found in Quick start. Additionally, the Quick startpage shows how to use Parasail's UI to monitor the progress of batch jobs.

To run the script below, execute the following commands:

pip3 install openai-batch
PARASAIL_API_KEY=<YOUR-API-KEY> python3 test_omnigen_batch.py 

#test_omnigen_batch.py

#pip install openai-batch
from openai_batch import Batch, providers
from PIL import Image
import os
import base64
import json
import io
from pathlib import Path


def extract_and_save_images(input_file: str, output_dir: str):
    """Extract base64 encoded images from batch output and save as JPG files."""
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    counter = 1
    with open(input_file, "r", encoding="utf-8") as f:
        for line in f:
            try:
                data = json.loads(line.strip())
                b64_data = data["response"]["body"]["data"][0]["b64_json"]
                image_bytes = base64.b64decode(b64_data)
                image = Image.open(io.BytesIO(image_bytes))
                output_path = os.path.join(output_dir, f"image-{counter}.jpg")
                image.convert("RGB").save(output_path, format="JPEG")
                print(f"Saved {output_path}")
                counter += 1
            except Exception as e:
                print(f"Error processing line {counter}: {e}")
                counter += 1
                continue


# Input and output directories
input_images_dir = Path("omnigen_input_images")
output_images_dir = Path("omnigen_output_images")

# Batch files
output_file = Path("output.jsonl")
error_file = Path("error.jsonl")
submission_file = Path("batch_submission.jsonl")

# Get all image files from input directory
image_extensions = {".jpg", ".jpeg", ".png"}
image_files = [
    f
    for f in input_images_dir.iterdir()
    if f.is_file() and f.suffix.lower() in image_extensions
]

if not image_files:
    print(f"No image files found in {input_images_dir}")
    print(f"Supported extensions: {', '.join(image_extensions)}")
    exit(1)

print(f"Found {len(image_files)} images to process")

# Create a batch with transfusion request
with Batch(
    submission_input_file=submission_file,
    output_file=output_file,
    error_file=error_file,
) as batch_obj:
    # Process each image and add to batch
    for image_path in image_files:
        print(f"Processing: {image_path.name}")

        # Read and encode the image as base64
        with open(image_path, "rb") as img_file:
            image_data = img_file.read()
            base64_image = base64.b64encode(image_data).decode("utf-8")

        # Add transfusion request to the batch
        batch_obj.add_to_batch(
            model="Shitao/OmniGen-v1",
            prompt="A man in a black shirt is reading a book. The man is the right man in <img><|image_1|></img>.",
            size="1024x1024",
            image=[base64_image],
            response_format="b64_json",
        )

    print(f"\nSubmitting batch with {len(image_files)} requests...")

    # Submit, wait for completion, and download results
    result, output_path, error_path = batch_obj.submit_wait_download()

    # Verify the batch completed successfully
    assert result.status == "completed", f"Batch failed with status: {result.status}"

# Verify output file exists
assert output_file.exists(), "Output file not created for transfusion"

# Extract and save generated images
extract_and_save_images(str(output_file), str(output_images_dir))