Image Generation and Editing
Parameter
Description
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))
Last updated