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_in = inputs.get("seed", -1)26        seed_input = int(seed_in) if seed_in not in (None, "") else -12728        if not prompt:29            raise ValueError("Prompt is required")3031        # Ernie supports a fixed set of (W, H) pairs. If the user provided32        # an explicit `resolution` override, parse it; otherwise compute33        # target dims from aspect_ratio + megapixel and snap to the34        # closest supported pair.35        ERNIE_SIZES = [36            (1024, 1024), (848, 1264), (1264, 848),37            (768, 1376), (1376, 768), (896, 1200), (1200, 896),38        ]3940        if resolution:41            try:42                w_str, h_str = resolution.lower().split("x")43                req_w, req_h = int(w_str), int(h_str)44            except Exception as e:45                raise ValueError(f"Invalid resolution '{resolution}': {e}")46        else:47            import math48            target_pixels = max(0.25, min(4.0, megapixel)) * 1_000_000.049            req_h = int(round(math.sqrt(target_pixels / aspect_ratio)))50            req_w = int(round(req_h * aspect_ratio))5152        # Snap to nearest Ernie-supported (W, H) by sum-of-deltas distance.53        width, height = min(54            ERNIE_SIZES,55            key=lambda wh: abs(wh[0] - req_w) + abs(wh[1] - req_h),56        )57        if (width, height) != (req_w, req_h):58            print(59                f"[snap image.generate_typographic] requested={req_w}x{req_h} "60                f"actual={width}x{height} (Ernie supported set)",61                file=sys.stderr,62            )6364        if seed_mode == "fixed" and seed_input >= 0:65            seed_value = seed_input66        else:67            seed_value = random.randint(0, 2**31 - 1)6869        print(70            f"Requesting Ernie generation: resolution={resolution}, "71            f"seed={seed_value}, use_pe={use_pe}",72            file=sys.stderr,73        )7475        result = Gais.image.generate_typographic(76            prompt=prompt,77            width=width,78            height=height,79            seed=seed_value,80            use_pe=use_pe,81        )8283        os.makedirs(OUTPUT_DIR, exist_ok=True)84        out_filename = "typographic_image.png"85        out_path = os.path.join(OUTPUT_DIR, out_filename)86        with open(out_path, "wb") as f:87            f.write(result.content)8889        seed_used = result.metadata.get("seed", str(seed_value))90        inference_time = result.metadata.get("inference_time_ms", "unknown")91        variant = result.metadata.get("variant", "unknown")92        print(93            f"Generated: seed={seed_used}, time={inference_time}ms, variant={variant}",94            file=sys.stderr,95        )9697        # Probe output dimensions and emit aspect_ratio + resolution so98        # downstream nodes can chain.99        try:100            from PIL import Image101            with Image.open(out_path) as _im:102                _w, _h = _im.size103        except Exception:104            _w, _h = 0, 0105106        output = {107            "image": out_filename,108            "aspect_ratio": round(_w / _h, 4) if _h else 0.0,109            "resolution": f"{_w}x{_h}",110        }111        print(json.dumps(output, indent=2))112113    except Exception as e:114        error_output = {115            "error": str(e),116            "errorType": type(e).__name__,117            "traceback": traceback.format_exc(),118        }119        print(json.dumps(error_output), file=sys.stderr)120        sys.exit(1)121122123if __name__ == "__main__":124    main()