Image Creation

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", {})17        service_defaults = execution_input.get("service_defaults", {})1819        prompt = inputs.get("prompt", "")20        aspect_ratio = float(inputs.get("aspect_ratio", 1.667))21        megapixel = float(inputs.get("megapixel", 1.0))2223        seed_mode = inputs.get("seed_mode", "random")24        seed_input = int(inputs.get("seed", -1))2526        if seed_mode == "fixed" and seed_input >= 0:27            seed_value = seed_input28        else:29            seed_value = random.randint(0, 2**31 - 1)3031        if not prompt:32            raise ValueError("Prompt is required")33        if not (0.25 <= aspect_ratio <= 4.0):34            raise ValueError(f"Aspect ratio must be between 0.25 and 4.0, got {aspect_ratio}")35        if not (0.25 <= megapixel <= 4.0):36            raise ValueError(f"Megapixel must be between 0.25 and 4.0, got {megapixel}")3738        print(39            f"Requesting generation: aspect_ratio={aspect_ratio}, megapixel={megapixel}, seed={seed_value}",40            file=sys.stderr,41        )4243        result = Gais.image.generate(44            prompt=prompt,45            aspect_ratio=aspect_ratio,46            megapixel=megapixel,47            seed_mode=seed_mode,48            seed=seed_value,49        )5051        os.makedirs(OUTPUT_DIR, exist_ok=True)52        out_path = os.path.join(OUTPUT_DIR, "generated_image.png")53        with open(out_path, "wb") as f:54            f.write(result.content)5556        seed_used = result.metadata.get("seed", "unknown")57        inference_time = result.metadata.get("inference_time_ms", "unknown")58        print(59            f"Generated: seed={seed_used}, time={inference_time}ms",60            file=sys.stderr,61        )6263        # Probe output dimensions and emit aspect_ratio + resolution so64        # downstream nodes can chain (Number / String ports respectively).65        try:66            from PIL import Image67            with Image.open(out_path) as _im:68                _w, _h = _im.size69        except Exception:70            _w, _h = 0, 07172        output = {73            "image": "generated_image.png",74            "aspect_ratio": round(_w / _h, 4) if _h else 0.0,75            "resolution": f"{_w}x{_h}",76        }77        print(json.dumps(output, indent=2))7879    except Exception as e:80        error_output = {81            "error": str(e),82            "errorType": type(e).__name__,83            "traceback": traceback.format_exc(),84        }85        print(json.dumps(error_output), file=sys.stderr)86        sys.exit(1)878889if __name__ == "__main__":90    main()