Image Creation

12import os3import sys4import json5import random6import traceback78from gais import Gais910OUTPUT_DIR = "/data/output"111213def main():14    try:15        input_json = sys.stdin.read()16        execution_input = json.loads(input_json)17        inputs = execution_input.get("inputs", {})18        service_defaults = execution_input.get("service_defaults", {})1920        prompt = inputs.get("prompt", "")21        aspect_ratio = float(inputs.get("aspect_ratio", 1.667) or 1.667)22        megapixel = float(inputs.get("megapixel", 1.0) or 1.0)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 seed_mode == "fixed" and seed_input >= 0:29            seed_value = seed_input30        else:31            seed_value = random.randint(0, 2**31 - 1)3233        if not prompt:34            raise ValueError("Prompt is required")35        if not (0.25 <= aspect_ratio <= 4.0):36            raise ValueError(f"Aspect ratio must be between 0.25 and 4.0, got {aspect_ratio}")37        if not (0.25 <= megapixel <= 4.0):38            raise ValueError(f"Megapixel must be between 0.25 and 4.0, got {megapixel}")3940        print(41            f"Requesting generation: aspect_ratio={aspect_ratio}, megapixel={megapixel}, seed={seed_value}",42            file=sys.stderr,43        )4445        result = Gais.image.generate(46            prompt=prompt,47            aspect_ratio=aspect_ratio,48            megapixel=megapixel,49            seed_mode=seed_mode,50            seed=seed_value,51        )5253        os.makedirs(OUTPUT_DIR, exist_ok=True)54        out_path = os.path.join(OUTPUT_DIR, "generated_image.png")55        with open(out_path, "wb") as f:56            f.write(result.content)5758        seed_used = result.metadata.get("seed", "unknown")59        inference_time = result.metadata.get("inference_time_ms", "unknown")60        print(61            f"Generated: seed={seed_used}, time={inference_time}ms",62            file=sys.stderr,63        )6465        # Probe output dimensions and emit aspect_ratio + resolution so66        # downstream nodes can chain (Number / String ports respectively).67        try:68            from PIL import Image69            with Image.open(out_path) as _im:70                _w, _h = _im.size71        except Exception:72            _w, _h = 0, 07374        output = {75            "image": "generated_image.png",76            "aspect_ratio": round(_w / _h, 4) if _h else 0.0,77            "resolution": f"{_w}x{_h}",78        }79        print(json.dumps(output, indent=2))8081    except Exception as e:82        error_output = {83            "error": str(e),84            "errorType": type(e).__name__,85            "traceback": traceback.format_exc(),86        }87        print(json.dumps(error_output), file=sys.stderr)88        sys.exit(1)899091if __name__ == "__main__":92    main()93