scripts/gen_image.py

"""Minimal text-to-image generation against a vLLM-Omni endpoint.

Run from the Cosmos3-Super-Text2Image-4Step repo root:

    python scripts/gen_image.py \
        --prompt-file assets/example_caption.json \
        --output-path scripts/output.png

Use --endpoint http://<server-host>:8000 if the vLLM-Omni server is not local.
"""

import argparse
import base64
import json
from pathlib import Path

import requests

# Fixed generation settings: square 768px image.
WIDTH = 768
HEIGHT = 768


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate one T2I sample.")
    parser.add_argument("--endpoint", default="http://localhost:8000", help="vLLM-Omni endpoint base URL.")
    parser.add_argument("--prompt-file", type=Path, default=Path("assets/example_caption.json"))
    parser.add_argument("--output-path", type=Path, default=Path("scripts/output.png"))
    args = parser.parse_args()

    json_prompt = json.loads(args.prompt_file.read_text(encoding="utf-8"))
    json_prompt["resolution"] = {"H": HEIGHT, "W": WIDTH}
    json_prompt["aspect_ratio"] = "1,1"

    request_body = {
        "prompt": json.dumps(json_prompt, ensure_ascii=False),
        "size": f"{WIDTH}x{HEIGHT}",
        "n": 1,
        "guidance_scale": 1.0,
        "negative_prompt": "",
        "seed": 1143,
        "extra_args": {
            "guardrails": True,
            "use_resolution_template": False,
        },
    }

    endpoint = args.endpoint.rstrip("/")
    response = requests.post(
        f"{endpoint}/v1/images/generations",
        json=request_body,
        headers={"Content-Type": "application/json"},
        timeout=(10, 600),
    )
    response.raise_for_status()

    response_json = response.json()
    b64_data = response_json["data"][0]["b64_json"]
    image_bytes = base64.b64decode(b64_data)

    args.output_path.parent.mkdir(parents=True, exist_ok=True)
    args.output_path.write_bytes(image_bytes)
    print(f"Saved image to {args.output_path} ({len(image_bytes) / (1024 * 1024):.1f} MB)")


if __name__ == "__main__":
    main()