File size: 3,847 Bytes
7b6142f | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | #!/usr/bin/env python3
"""Enhance a prompt with the MiniMax H3-Image-Context-IR API.
export MINIMAX_API_BASE=https://api.minimax.cn # CN; Global: https://api.minimax.io
export MINIMAX_API_KEY=<API key>
python3 context_ir.py t2i # instruction.t2i.txt
python3 context_ir.py edit # instruction.edit.txt with reference1.png
Sends the instruction (and, for edit, the reference images) to POST /v1/h3_image_context_ir
and writes the enhanced prompt to revised_prompt.txt. Unless --size is given, the API also
picks the aspect ratio; the matching size of about 16 megapixels is written to
revised_size.txt. Pass the files to t2i.py / edit.py with --prompt and --size. Uses only the
Python standard library.
"""
import argparse
import base64
import json
import math
import mimetypes
import os
import urllib.request
from pathlib import Path
# absolute(), not resolve(): in the Hugging Face cache this file is a symlink into blobs/.
HERE = Path(__file__).absolute().parent
AREA = 4096 * 4096
def data_uri(path: Path) -> str:
content_type = mimetypes.guess_type(path.name)[0]
if content_type is None:
raise ValueError(f"unknown image type: {path}")
return f"data:{content_type};base64,{base64.b64encode(path.read_bytes()).decode()}"
def size_for(aspect_ratio: str) -> str:
"""About 16 megapixels at the ratio, both sides multiples of 32 (matches the README Sizes table)."""
x, y = (int(v) for v in aspect_ratio.split(":"))
width = round(math.sqrt(AREA * x / y) / 32) * 32
height = round(math.sqrt(AREA * y / x) / 32) * 32
return f"{width}x{height}"
def post(path: str, body: dict) -> dict:
request = urllib.request.Request(
f"{os.environ['MINIMAX_API_BASE'].rstrip('/')}{path}",
data=json.dumps(body).encode(),
headers={"Authorization": f"Bearer {os.environ['MINIMAX_API_KEY']}", "Content-Type": "application/json"},
)
with urllib.request.urlopen(request, timeout=600) as response:
return json.load(response)
def request_body(task: str, instruction: Path, images: list, size) -> dict:
body = {"model": "MiniMax-H3-Image", "prompt": instruction.read_text().strip()}
if task == "edit":
body["images"] = [{"image_url": data_uri(path)} for path in images or [HERE / "reference1.png"]]
if size:
body["size"] = size
return body
def add_task_arguments(parser: argparse.ArgumentParser) -> None:
parser.add_argument("task", choices=["t2i", "edit"])
parser.add_argument("--instruction", type=Path, help="default: instruction.<task>.txt")
parser.add_argument("--image", action="append", type=Path,
help="edit only: reference image; repeat for <Picture 2>, ... (default: reference1.png)")
parser.add_argument("--size", help="WIDTHxHEIGHT of the image you will generate (default: chosen by the API)")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
add_task_arguments(parser)
parser.add_argument("--out", default="revised_prompt.txt")
parser.add_argument("--size-out", default="revised_size.txt", help="written when the API picks the aspect ratio")
args = parser.parse_args()
instruction = args.instruction or HERE / f"instruction.{args.task}.txt"
result = post("/v1/h3_image_context_ir", request_body(args.task, instruction, args.image, args.size))
Path(args.out).write_text(result["revised_prompt"])
summary = {"request_id": result.get("request_id"), "usage": result.get("usage")}
if "aspect_ratio" in result:
summary["aspect_ratio"] = result["aspect_ratio"]
summary["size"] = size_for(result["aspect_ratio"])
Path(args.size_out).write_text(summary["size"])
print(json.dumps(summary))
if __name__ == "__main__":
main()
|