Typographic Image

1import os2import sys3import json4import random5import traceback67from gais import Gais89OUTPUT_DIR = "/data/output"101112def main():13    try:14        input_json = sys.stdin.read()15        execution_input = json.loads(input_json)16        inputs = execution_input.get("inputs", {})1718        prompt = inputs.get("prompt", "")19        aspect_ratio = float(inputs.get("aspect_ratio", 1.7778) or 1.7778)20        megapixel = float(inputs.get("megapixel", 1.0) or 1.0)21        resolution = (inputs.get("resolution") or "").strip()22        use_pe = bool(inputs.get("use_prompt_enhancer", True))2324        seed_mode = inputs.get("seed_mode", "random")25        seed_input = int(inputs.get("seed", -1))2627        if not prompt:28            raise ValueError("Prompt is required")2930        # Ernie supports a fixed set of (W, H) pairs. If the user provided31        # an explicit `resolution` override, parse it; otherwise compute32        # target dims from aspect_ratio + megapixel and snap to the33        # closest supported pair.34        ERNIE_SIZES = [35            (1024, 1024), (848, 1264), (1264, 848),36            (768, 1376), (1376, 768), (896, 1200), (1200, 896),37        ]3839        if resolution:40            try:41                w_str, h_str = resolution.lower().split("x")42                req_w, req_h = int(w_str), int(h_str)43            except Exception as e:44                raise ValueError(f"Invalid resolution '{resolution}': {e}")45        else:46            import math47            target_pixels = max(0.25, min(4.0, megapixel)) * 1_000_000.048            req_h = int(round(math.sqrt(target_pixels / aspect_ratio)))49            req_w = int(round(req_h * aspect_ratio))5051        # Snap to nearest Ernie-supported (W, H) by sum-of-deltas distance.52        width, height = min(53            ERNIE_SIZES,54            key=lambda wh: abs(wh[0] - req_w) + abs(wh[1] - req_h),55        )56        if (width, height) != (req_w, req_h):57            print(58                f"[snap image.generate_typographic] requested={req_w}x{req_h} "59                f"actual={width}x{height} (Ernie supported set)",60                file=sys.stderr,61            )6263        if seed_mode == "fixed" and seed_input >= 0:64            seed_value = seed_input65        else:66            seed_value = random.randint(0, 2**31 - 1)6768        print(69            f"Requesting Ernie generation: resolution={resolution}, "70            f"seed={seed_value}, use_pe={use_pe}",71            file=sys.stderr,72        )7374        result = Gais.image.generate_typographic(75            prompt=prompt,76            width=width,77            height=height,78            seed=seed_value,79            use_pe=use_pe,80        )8182        os.makedirs(OUTPUT_DIR, exist_ok=True)83        out_filename = "typographic_image.png"84        out_path = os.path.join(OUTPUT_DIR, out_filename)85        with open(out_path, "wb") as f:86            f.write(result.content)8788        seed_used = result.metadata.get("seed", str(seed_value))89        inference_time = result.metadata.get("inference_time_ms", "unknown")90        variant = result.metadata.get("variant", "unknown")91        print(92            f"Generated: seed={seed_used}, time={inference_time}ms, variant={variant}",93            file=sys.stderr,94        )9596        # Probe output dimensions and emit aspect_ratio + resolution so97        # downstream nodes can chain.98        try:99            from PIL import Image100            with Image.open(out_path) as _im:101                _w, _h = _im.size102        except Exception:103            _w, _h = 0, 0104105        output = {106            "image": out_filename,107            "aspect_ratio": round(_w / _h, 4) if _h else 0.0,108            "resolution": f"{_w}x{_h}",109        }110        print(json.dumps(output, indent=2))111112    except Exception as e:113        error_output = {114            "error": str(e),115            "errorType": type(e).__name__,116            "traceback": traceback.format_exc(),117        }118        print(json.dumps(error_output), file=sys.stderr)119        sys.exit(1)120121122if __name__ == "__main__":123    main()