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))
---
# Agent Instructions: Querying This Documentation
If you need additional information that is not directly available in this page, you can query the documentation dynamically by asking a question.
Perform an HTTP GET request on the current page URL with the `ask` query parameter:
```
GET https://docs.parasail.io/parasail-docs/image-generation-and-editing.md?ask=
```
The question should be specific, self-contained, and written in natural language.
The response will contain a direct answer to the question and relevant excerpts and sources from the documentation.
Use this mechanism when the answer is not explicitly present in the current page, you need clarification or additional context, or you want to retrieve related documentation sections.