File size: 2,071 Bytes
6319a87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """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()
|