Image Edit

1import os2import sys3import json4import random5import traceback67from gais import Gais89INPUT_DIR = "/data/input"10OUTPUT_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        images = inputs.get("images", [])21        prompt = inputs.get("prompt", "")22        aspect_ratio = float(inputs.get("aspect_ratio", 1.667))23        megapixel = float(inputs.get("megapixel", 1.0))2425        # Seed mode handling26        seed_mode = inputs.get("seed_mode", "random")27        seed_input = int(inputs.get("seed", -1))2829        if seed_mode == "fixed" and seed_input >= 0:30            seed_value = seed_input31        else:32            seed_value = random.randint(0, 2**31 - 1)3334        # Normalize images to a list (single edge gives a string, array edge gives a list)35        if isinstance(images, str):36            images = [images]3738        if not prompt:39            raise ValueError("Prompt is required")40        if not (0.25 <= aspect_ratio <= 4.0):41            raise ValueError(f"Aspect ratio must be between 0.25 and 4.0, got {aspect_ratio}")42        if not (0.25 <= megapixel <= 4.0):43            raise ValueError(f"Megapixel must be between 0.25 and 4.0, got {megapixel}")44        if not images or len(images) == 0:45            raise ValueError("At least one input image is required")46        if len(images) > 2:47            raise ValueError("Maximum of 2 input images supported")4849        # Validate all input images exist and collect paths50        image_paths = []51        for img_filename in images:52            local_path = os.path.join(INPUT_DIR, img_filename)53            if not os.path.exists(local_path):54                raise FileNotFoundError(f"Input image not found: {local_path}")55            image_paths.append(local_path)5657        os.makedirs(OUTPUT_DIR, exist_ok=True)5859        print(60            f"Requesting edit: images={len(image_paths)}, aspect_ratio={aspect_ratio}, megapixel={megapixel}, seed={seed_value}",61            file=sys.stderr,62        )6364        result = Gais.image.edit(65            images=image_paths,66            prompt=prompt,67            aspect_ratio=aspect_ratio,68            megapixel=megapixel,69            seed_mode=seed_mode,70            seed=seed_value,71        )7273        # Save result74        out_filename = "edited_image.png"75        out_path = os.path.join(OUTPUT_DIR, out_filename)76        with open(out_path, "wb") as f:77            f.write(result.content)7879        seed_used = result.metadata.get("seed", str(seed_value))80        inference_time = result.metadata.get("inference_time_ms", "unknown")81        output_size = result.metadata.get("output_size", "unknown")82        print(83            f"Edited: seed={seed_used}, time={inference_time}ms, size={output_size}, images={len(image_paths)}",84            file=sys.stderr,85        )8687        # Probe output dimensions and emit aspect_ratio + resolution so88        # downstream nodes can chain.89        try:90            from PIL import Image91            with Image.open(out_path) as _im:92                _w, _h = _im.size93        except Exception:94            _w, _h = 0, 09596        # Flat output — keys match OUTPUT_SCHEMA97        output = {98            "image": out_filename,99            "aspect_ratio": round(_w / _h, 4) if _h else 0.0,100            "resolution": f"{_w}x{_h}",101        }102        print(json.dumps(output, indent=2))103104    except Exception as e:105        error_output = {106            "error": str(e),107            "errorType": type(e).__name__,108            "traceback": traceback.format_exc(),109        }110        print(json.dumps(error_output), file=sys.stderr)111        sys.exit(1)112113114if __name__ == "__main__":115    main()