Image Edit

12import os3import sys4import json5import random6import tempfile7import mimetypes8import traceback910from gais import Gais1112INPUT_DIR = "/data/input"13OUTPUT_DIR = "/data/output"1415# aspect_ratio sentinel for "Match input image" (same value as the module-level16# MATCH_INPUT_ASPECT in image_edit_native.py).17MATCH_INPUT_ASPECT = -11819# EXIF Orientation tag. Values 2-8 mean the stored pixels must be flipped or20# rotated to display upright; 1 (or no tag) means they already are.21EXIF_ORIENTATION = 0x011222TRANSPOSING_ORIENTATIONS = (2, 3, 4, 5, 6, 7, 8)232425def _is_match_input(raw):26    """True for the "Match input image" sentinel: -1 as an int, float or string."""27    try:28        return float(raw) == MATCH_INPUT_ASPECT29    except (TypeError, ValueError):30        return False313233def _check_aspect_ratio(aspect_ratio):34    if not (0.25 <= aspect_ratio <= 4.0):35        raise ValueError(f"Aspect ratio must be between 0.25 and 4.0, got {aspect_ratio}")363738def _save_upright_copy(src_path, upright, src_format):39    """Write the EXIF-transposed pixels to a new file and return its path.4041    /data/input is read-only in both executors (Docker mounts it :ro, the warm42    pool chmods it 0555), so the copy goes to the run's temp dir; the GAIS SDK43    stages input files from any path. Same format and extension as the source,44    so the backend guesses the same mime. The copy carries no EXIF, so no45    backend can apply the orientation a second time.46    """47    scratch = tempfile.mkdtemp(prefix="image_edit_upright_")48    stem, ext = os.path.splitext(os.path.basename(src_path))49    save_kwargs = {}50    if upright.info.get("icc_profile"):51        save_kwargs["icc_profile"] = upright.info["icc_profile"]52    fmt = (src_format or "").upper()53    if fmt in ("JPEG", "MPO"):54        # Pillow opens many phone JPEGs as MPO; send a plain JPEG.55        save_kwargs.update(format="JPEG", quality=95)56        ext = ext or ".jpg"57    else:58        save_kwargs["format"] = fmt59    dst = os.path.join(scratch, f"{stem}_upright{ext}")60    upright.save(dst, **save_kwargs)61    return dst626364def _resolve_match_input(image_path):65    """Measure the first reference image for "Match input image".6667    Returns (aspect_ratio, path to send as the first image). The ratio is68    width / height after applying the EXIF orientation. Cloud backends (xAI)69    receive the raw file bytes, so when the orientation flips or rotates the70    pixels the backend gets an upright copy instead of the original.71    """72    from PIL import Image, ImageOps7374    upright = None75    try:76        with Image.open(image_path) as src:77            src_format = src.format78            # For a PNG, getexif() decodes the pixels to reach an eXIf chunk79            # stored after the image data, a "Raw profile type exif" text80            # chunk or XMP. Accepted: a header-only shortcut misses all three81            # and would send a sideways image.82            orientation = src.getexif().get(EXIF_ORIENTATION, 1)83            size = src.size84            if orientation in TRANSPOSING_ORIENTATIONS:85                # A new image holding the transposed pixels; outlives src.86                upright = ImageOps.exif_transpose(src)87                size = upright.size88    except Image.DecompressionBombError:89        raise ValueError(90            "Image is too large to read (over ~179 megapixels); downscale it and try again"91        ) from None92    except Exception as exc:93        ext = os.path.splitext(image_path)[1].lower()94        mime = mimetypes.guess_type(image_path)[0] or "unknown format"95        # gstack-shortcut(dec-bf78833f): HEIC needs manual conversion, upgrade when the first real HEIC upload arrives96        raise ValueError(f"Can't read image ({ext or mime}); convert to JPG or PNG") from exc9798    width, height = size99    send_path = image_path100    if upright is not None:101        send_path = _save_upright_copy(image_path, upright, src_format)102        print(103            f"Match input: EXIF orientation {orientation}, sending an upright copy ({width}x{height})",104            file=sys.stderr,105        )106    print(f"Match input: {width}x{height} -> aspect_ratio={width / height:.4f}", file=sys.stderr)107    return width / height, send_path108109110def main():111    try:112        input_json = sys.stdin.read()113        execution_input = json.loads(input_json)114        inputs = execution_input.get("inputs", {})115        service_defaults = execution_input.get("service_defaults", {})116117        images = inputs.get("images", [])118        prompt = inputs.get("prompt", "")119        # Checked before the fallback below: "Match input image" is measured120        # from the first image once its path is validated. Any other value121        # (absent, falsy, numeric) behaves exactly as before.122        match_input = _is_match_input(inputs.get("aspect_ratio"))123        aspect_ratio = None124        if not match_input:125            aspect_ratio = float(inputs.get("aspect_ratio", 1.667) or 1.667)126        megapixel = float(inputs.get("megapixel", 1.0) or 1.0)127128        # Seed mode handling129        seed_mode = inputs.get("seed_mode", "random")130        seed_in = inputs.get("seed", -1)131        seed_input = int(seed_in) if seed_in not in (None, "") else -1132133        if seed_mode == "fixed" and seed_input >= 0:134            seed_value = seed_input135        else:136            seed_value = random.randint(0, 2**31 - 1)137138        # Normalize images to a list (single edge gives a string, array edge gives a list)139        if isinstance(images, str):140            images = [images]141142        if not prompt:143            raise ValueError("Prompt is required")144        if not match_input:145            _check_aspect_ratio(aspect_ratio)146        if not (0.25 <= megapixel <= 4.0):147            raise ValueError(f"Megapixel must be between 0.25 and 4.0, got {megapixel}")148        if not images or len(images) == 0:149            raise ValueError("At least one input image is required")150        if len(images) > 10:151            raise ValueError("Maximum of 10 input images supported")152153        # Validate all input images exist and collect paths154        image_paths = []155        for img_filename in images:156            local_path = os.path.join(INPUT_DIR, img_filename)157            if not os.path.exists(local_path):158                raise FileNotFoundError(f"Input image not found: {local_path}")159            image_paths.append(local_path)160161        if match_input:162            aspect_ratio, image_paths[0] = _resolve_match_input(image_paths[0])163            if not (0.25 <= aspect_ratio <= 4.0):164                raise ValueError(165                    f'"Match input image": the first image\'s aspect ratio is {aspect_ratio:.2f}, '166                    "outside the 0.25-4.0 Image Edit supports; pick a fixed aspect ratio instead"167                )168169        os.makedirs(OUTPUT_DIR, exist_ok=True)170171        print(172            f"Requesting edit: images={len(image_paths)}, aspect_ratio={aspect_ratio}, megapixel={megapixel}, seed={seed_value}",173            file=sys.stderr,174        )175176        result = Gais.image.edit(177            images=image_paths,178            prompt=prompt,179            aspect_ratio=aspect_ratio,180            megapixel=megapixel,181            seed_mode=seed_mode,182            seed=seed_value,183        )184185        # Save result186        out_filename = "edited_image.png"187        out_path = os.path.join(OUTPUT_DIR, out_filename)188        with open(out_path, "wb") as f:189            f.write(result.content)190191        seed_used = result.metadata.get("seed", str(seed_value))192        inference_time = result.metadata.get("inference_time_ms", "unknown")193        output_size = result.metadata.get("output_size", "unknown")194        print(195            f"Edited: seed={seed_used}, time={inference_time}ms, size={output_size}, images={len(image_paths)}",196            file=sys.stderr,197        )198199        # Probe output dimensions and emit aspect_ratio + resolution so200        # downstream nodes can chain.201        try:202            from PIL import Image203            with Image.open(out_path) as _im:204                _w, _h = _im.size205        except Exception:206            _w, _h = 0, 0207208        # Flat output — keys match OUTPUT_SCHEMA209        output = {210            "image": out_filename,211            "aspect_ratio": round(_w / _h, 4) if _h else 0.0,212            "resolution": f"{_w}x{_h}",213        }214        print(json.dumps(output, indent=2))215216    except Exception as e:217        error_output = {218            "error": str(e),219            "errorType": type(e).__name__,220            "traceback": traceback.format_exc(),221        }222        print(json.dumps(error_output), file=sys.stderr)223        sys.exit(1)224225226if __name__ == "__main__":227    main()228