| |
| """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 |
|
|
| |
| 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() |
|
|
