Image Edit Legacy

1import os2import sys3import json4import time5import random6import traceback7from fractions import Fraction8from datetime import datetime910try:11    import requests12except ImportError:13    import subprocess14    subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])15    import requests1617COMFYUI_API_URL = os.getenv("COMFYUI_API_URL", "http://192.168.1.39:8188")18INPUT_DIR = "/data/input"19OUTPUT_DIR = "/data/output"2021# ---------- Single-image workflow (EM_Edit_Image_single_v1) ----------22WORKFLOW_SINGLE = {23    "9": {24        "inputs": {"filename_prefix": "Flux2-Klein", "images": ["75:65", 0]},25        "class_type": "SaveImage",26        "_meta": {"title": "Save Image"},27    },28    "120": {29        "inputs": {30            "megapixel": "1.0",31            "aspect_ratio": "9:16 (Slim Vertical)",32            "divisible_by": "64",33            "custom_ratio": True,34            "custom_aspect_ratio": "1:1",35        },36        "class_type": "FluxResolutionNode",37        "_meta": {"title": "Flux Resolution Calc"},38    },39    "132": {40        "inputs": {"image": "", "custom_width": 0, "custom_height": 0},41        "class_type": "VHS_LoadImagePath",42        "_meta": {"title": "Load Image (Path)"},43    },44    "75:61": {45        "inputs": {"sampler_name": "euler"},46        "class_type": "KSamplerSelect",47        "_meta": {"title": "KSamplerSelect"},48    },49    "75:64": {50        "inputs": {51            "noise": ["75:73", 0],52            "guider": ["75:63", 0],53            "sampler": ["75:61", 0],54            "sigmas": ["75:62", 0],55            "latent_image": ["75:66", 0],56        },57        "class_type": "SamplerCustomAdvanced",58        "_meta": {"title": "SamplerCustomAdvanced"},59    },60    "75:65": {61        "inputs": {"samples": ["75:64", 0], "vae": ["75:72", 0]},62        "class_type": "VAEDecode",63        "_meta": {"title": "VAE Decode"},64    },65    "75:73": {66        "inputs": {"noise_seed": 0},67        "class_type": "RandomNoise",68        "_meta": {"title": "RandomNoise"},69    },70    "75:70": {71        "inputs": {72            "unet_name": "flux-2-klein-9b-fp8.safetensors",73            "weight_dtype": "default",74        },75        "class_type": "UNETLoader",76        "_meta": {"title": "Load Diffusion Model"},77    },78    "75:71": {79        "inputs": {80            "clip_name": "qwen_3_8b_fp8mixed.safetensors",81            "type": "flux2",82            "device": "default",83        },84        "class_type": "CLIPLoader",85        "_meta": {"title": "Load CLIP"},86    },87    "75:72": {88        "inputs": {"vae_name": "flux2-vae.safetensors"},89        "class_type": "VAELoader",90        "_meta": {"title": "Load VAE"},91    },92    "75:63": {93        "inputs": {94            "cfg": 1,95            "model": ["75:70", 0],96            "positive": ["75:79:77", 0],97            "negative": ["75:79:76", 0],98        },99        "class_type": "CFGGuider",100        "_meta": {"title": "CFGGuider"},101    },102    "75:62": {103        "inputs": {104            "steps": 4,105            "width": ["75:81", 0],106            "height": ["75:81", 1],107        },108        "class_type": "Flux2Scheduler",109        "_meta": {"title": "Flux2Scheduler"},110    },111    "75:74": {112        "inputs": {"text": "", "clip": ["75:71", 0]},113        "class_type": "CLIPTextEncode",114        "_meta": {"title": "CLIP Text Encode (Positive Prompt)"},115    },116    "75:82": {117        "inputs": {"conditioning": ["75:74", 0]},118        "class_type": "ConditioningZeroOut",119        "_meta": {"title": "ConditioningZeroOut"},120    },121    "75:80": {122        "inputs": {123            "upscale_method": "nearest-exact",124            "megapixels": 1,125            "resolution_steps": 1,126            "image": ["132", 0],127        },128        "class_type": "ImageScaleToTotalPixels",129        "_meta": {"title": "ImageScaleToTotalPixels"},130    },131    "75:81": {132        "inputs": {"image": ["75:80", 0]},133        "class_type": "GetImageSize",134        "_meta": {"title": "Get Image Size"},135    },136    "75:66": {137        "inputs": {138            "width": ["75:84", 0],139            "height": ["75:83", 0],140            "batch_size": 1,141        },142        "class_type": "EmptyFlux2LatentImage",143        "_meta": {"title": "Empty Flux 2 Latent"},144    },145    "75:79:76": {146        "inputs": {147            "conditioning": ["75:82", 0],148            "latent": ["75:79:78", 0],149        },150        "class_type": "ReferenceLatent",151        "_meta": {"title": "ReferenceLatent"},152    },153    "75:79:78": {154        "inputs": {"pixels": ["75:80", 0], "vae": ["75:72", 0]},155        "class_type": "VAEEncode",156        "_meta": {"title": "VAE Encode"},157    },158    "75:79:77": {159        "inputs": {160            "conditioning": ["75:74", 0],161            "latent": ["75:79:78", 0],162        },163        "class_type": "ReferenceLatent",164        "_meta": {"title": "ReferenceLatent"},165    },166    "75:84": {167        "inputs": {"value": ["120", 0]},168        "class_type": "PrimitiveInt",169        "_meta": {"title": "Width"},170    },171    "75:83": {172        "inputs": {"value": ["120", 1]},173        "class_type": "PrimitiveInt",174        "_meta": {"title": "Height"},175    },176}177178# ---------- Multi-image workflow (EM_Edit_Image_multiple_v1) ----------179WORKFLOW_MULTIPLE = {180    "94": {181        "inputs": {"filename_prefix": "Flux2-Klein", "images": ["123:65", 0]},182        "class_type": "SaveImage",183        "_meta": {"title": "Save Image"},184    },185    "124": {186        "inputs": {187            "megapixel": "1.0",188            "aspect_ratio": "9:16 (Slim Vertical)",189            "divisible_by": "64",190            "custom_ratio": True,191            "custom_aspect_ratio": "1:1",192        },193        "class_type": "FluxResolutionNode",194        "_meta": {"title": "Flux Resolution Calc"},195    },196    "134": {197        "inputs": {"image": "", "custom_width": 0, "custom_height": 0},198        "class_type": "VHS_LoadImagePath",199        "_meta": {"title": "Load Image (Path)"},200    },201    "135": {202        "inputs": {"image": "", "custom_width": 0, "custom_height": 0},203        "class_type": "VHS_LoadImagePath",204        "_meta": {"title": "Load Image (Path) (2)"},205    },206    "123:61": {207        "inputs": {"sampler_name": "euler"},208        "class_type": "KSamplerSelect",209        "_meta": {"title": "KSamplerSelect"},210    },211    "123:62": {212        "inputs": {213            "steps": 4,214            "width": ["123:81", 0],215            "height": ["123:81", 1],216        },217        "class_type": "Flux2Scheduler",218        "_meta": {"title": "Flux2Scheduler"},219    },220    "123:63": {221        "inputs": {222            "cfg": 1,223            "model": ["123:70", 0],224            "positive": ["123:84:77", 0],225            "negative": ["123:84:76", 0],226        },227        "class_type": "CFGGuider",228        "_meta": {"title": "CFGGuider"},229    },230    "123:64": {231        "inputs": {232            "noise": ["123:73", 0],233            "guider": ["123:63", 0],234            "sampler": ["123:61", 0],235            "sigmas": ["123:62", 0],236            "latent_image": ["123:66", 0],237        },238        "class_type": "SamplerCustomAdvanced",239        "_meta": {"title": "SamplerCustomAdvanced"},240    },241    "123:65": {242        "inputs": {"samples": ["123:64", 0], "vae": ["123:72", 0]},243        "class_type": "VAEDecode",244        "_meta": {"title": "VAE Decode"},245    },246    "123:73": {247        "inputs": {"noise_seed": 0},248        "class_type": "RandomNoise",249        "_meta": {"title": "RandomNoise"},250    },251    "123:70": {252        "inputs": {253            "unet_name": "flux-2-klein-9b-fp8.safetensors",254            "weight_dtype": "default",255        },256        "class_type": "UNETLoader",257        "_meta": {"title": "Load Diffusion Model"},258    },259    "123:71": {260        "inputs": {261            "clip_name": "qwen_3_8b_fp8mixed.safetensors",262            "type": "flux2",263            "device": "default",264        },265        "class_type": "CLIPLoader",266        "_meta": {"title": "Load CLIP"},267    },268    "123:74": {269        "inputs": {"text": "", "clip": ["123:71", 0]},270        "class_type": "CLIPTextEncode",271        "_meta": {"title": "CLIP Text Encode (Positive Prompt)"},272    },273    "123:72": {274        "inputs": {"vae_name": "flux2-vae.safetensors"},275        "class_type": "VAELoader",276        "_meta": {"title": "Load VAE"},277    },278    "123:80": {279        "inputs": {280            "upscale_method": "nearest-exact",281            "megapixels": 1,282            "resolution_steps": 1,283            "image": ["134", 0],284        },285        "class_type": "ImageScaleToTotalPixels",286        "_meta": {"title": "ImageScaleToTotalPixels"},287    },288    "123:85": {289        "inputs": {290            "upscale_method": "nearest-exact",291            "megapixels": 1,292            "resolution_steps": 1,293            "image": ["135", 0],294        },295        "class_type": "ImageScaleToTotalPixels",296        "_meta": {"title": "ImageScaleToTotalPixels"},297    },298    "123:79:76": {299        "inputs": {300            "conditioning": ["123:86", 0],301            "latent": ["123:79:78", 0],302        },303        "class_type": "ReferenceLatent",304        "_meta": {"title": "ReferenceLatent"},305    },306    "123:79:78": {307        "inputs": {"pixels": ["123:80", 0], "vae": ["123:72", 0]},308        "class_type": "VAEEncode",309        "_meta": {"title": "VAE Encode"},310    },311    "123:79:77": {312        "inputs": {313            "conditioning": ["123:74", 0],314            "latent": ["123:79:78", 0],315        },316        "class_type": "ReferenceLatent",317        "_meta": {"title": "ReferenceLatent"},318    },319    "123:84:76": {320        "inputs": {321            "conditioning": ["123:79:76", 0],322            "latent": ["123:84:78", 0],323        },324        "class_type": "ReferenceLatent",325        "_meta": {"title": "ReferenceLatent"},326    },327    "123:84:78": {328        "inputs": {"pixels": ["123:85", 0], "vae": ["123:72", 0]},329        "class_type": "VAEEncode",330        "_meta": {"title": "VAE Encode"},331    },332    "123:84:77": {333        "inputs": {334            "conditioning": ["123:79:77", 0],335            "latent": ["123:84:78", 0],336        },337        "class_type": "ReferenceLatent",338        "_meta": {"title": "ReferenceLatent"},339    },340    "123:86": {341        "inputs": {"conditioning": ["123:74", 0]},342        "class_type": "ConditioningZeroOut",343        "_meta": {"title": "ConditioningZeroOut"},344    },345    "123:66": {346        "inputs": {347            "width": ["123:88", 0],348            "height": ["123:87", 0],349            "batch_size": 1,350        },351        "class_type": "EmptyFlux2LatentImage",352        "_meta": {"title": "Empty Flux 2 Latent"},353    },354    "123:81": {355        "inputs": {"image": ["123:80", 0]},356        "class_type": "GetImageSize",357        "_meta": {"title": "Get Image Size"},358    },359    "123:88": {360        "inputs": {"value": ["124", 0]},361        "class_type": "PrimitiveInt",362        "_meta": {"title": "Width"},363    },364    "123:87": {365        "inputs": {"value": ["124", 1]},366        "class_type": "PrimitiveInt",367        "_meta": {"title": "Height"},368    },369}370371372def aspect_ratio_to_string(ratio: float) -> str:373    """Convert a float w/h ratio to a W:H string."""374    frac = Fraction(ratio).limit_denominator(64)375    return f"{frac.numerator}:{frac.denominator}"376377378def upload_image_to_comfyui(local_path: str) -> str:379    """Upload a local image to ComfyUI and return the uploaded filename."""380    with open(local_path, "rb") as f:381        resp = requests.post(382            f"{COMFYUI_API_URL}/upload/image",383            files={"image": (os.path.basename(local_path), f, "image/png")},384            timeout=30,385        )386    resp.raise_for_status()387    data = resp.json()388    return data["name"]389390391def build_single_workflow(prompt: str, aspect_ratio: float, image_name: str) -> dict:392    """Build single-image edit workflow."""393    import copy394    wf = copy.deepcopy(WORKFLOW_SINGLE)395396    # Set prompt397    wf["75:74"]["inputs"]["text"] = prompt398399    # Set aspect ratio400    ratio_str = aspect_ratio_to_string(aspect_ratio)401    wf["120"]["inputs"]["custom_ratio"] = True402    wf["120"]["inputs"]["custom_aspect_ratio"] = ratio_str403404    # Set input image — swap VHS_LoadImagePath to LoadImage405    wf["132"] = {406        "inputs": {"image": image_name},407        "class_type": "LoadImage",408        "_meta": {"title": "Load Image"},409    }410411    # Randomize noise seed412    wf["75:73"]["inputs"]["noise_seed"] = random.randint(0, 2**31 - 1)413414    return wf415416417def build_multiple_workflow(prompt: str, aspect_ratio: float, image_name_1: str, image_name_2: str) -> dict:418    """Build multi-image edit workflow."""419    import copy420    wf = copy.deepcopy(WORKFLOW_MULTIPLE)421422    # Set prompt423    wf["123:74"]["inputs"]["text"] = prompt424425    # Set aspect ratio426    ratio_str = aspect_ratio_to_string(aspect_ratio)427    wf["124"]["inputs"]["custom_ratio"] = True428    wf["124"]["inputs"]["custom_aspect_ratio"] = ratio_str429430    # Set input images — swap VHS_LoadImagePath to LoadImage431    wf["134"] = {432        "inputs": {"image": image_name_1},433        "class_type": "LoadImage",434        "_meta": {"title": "Load Image"},435    }436    wf["135"] = {437        "inputs": {"image": image_name_2},438        "class_type": "LoadImage",439        "_meta": {"title": "Load Image (2)"},440    }441442    # Randomize noise seed443    wf["123:73"]["inputs"]["noise_seed"] = random.randint(0, 2**31 - 1)444445    return wf446447448def submit_prompt(workflow: dict) -> str:449    """Submit workflow to ComfyUI and return prompt_id."""450    resp = requests.post(451        f"{COMFYUI_API_URL}/prompt",452        json={"prompt": workflow},453        timeout=30,454    )455    if resp.status_code != 200:456        try:457            error_detail = resp.json()458        except Exception:459            error_detail = resp.text460        raise RuntimeError(461            f"ComfyUI /prompt returned {resp.status_code}: {json.dumps(error_detail, indent=2) if isinstance(error_detail, dict) else error_detail}"462        )463    data = resp.json()464465    # ComfyUI returns 200 even when nodes have validation errors466    node_errors = data.get("node_errors", {})467    if node_errors:468        raise RuntimeError(469            f"ComfyUI workflow has node errors: {json.dumps(node_errors, indent=2)}"470        )471472    return data["prompt_id"]473474475def wait_for_result(prompt_id: str, timeout: int = 600, poll_interval: int = 2) -> dict:476    """Poll ComfyUI history until the prompt completes with outputs."""477    deadline = time.time() + timeout478    empty_complete_retries = 0479    max_empty_retries = 3  # grace period for output serialization lag480481    while time.time() < deadline:482        resp = requests.get(483            f"{COMFYUI_API_URL}/history/{prompt_id}",484            timeout=10,485        )486        resp.raise_for_status()487        history = resp.json()488489        if prompt_id in history:490            prompt_data = history[prompt_id]491            status = prompt_data.get("status", {})492493            if status.get("status_str") == "error":494                messages = status.get("messages", [])495                raise RuntimeError(496                    f"ComfyUI prompt failed: {json.dumps(messages, indent=2)}"497                )498499            if status.get("completed", False):500                if prompt_data.get("outputs"):501                    return prompt_data502503                # Completed but no outputs — retry briefly for race condition504                empty_complete_retries += 1505                if empty_complete_retries >= max_empty_retries:506                    raise RuntimeError(507                        f"ComfyUI prompt completed but produced no outputs. "508                        f"This usually means a node failed silently (missing custom node or model). "509                        f"Status: {json.dumps(status, indent=2)}"510                    )511512        time.sleep(poll_interval)513514    raise TimeoutError(f"ComfyUI prompt {prompt_id} did not complete within {timeout}s")515516517def download_output_image(prompt_data: dict, output_dir: str) -> str:518    """Download the generated image from ComfyUI."""519    outputs = prompt_data.get("outputs", {})520    for node_id, node_output in outputs.items():521        images = node_output.get("images", [])522        if images:523            img_info = images[0]524            filename = img_info["filename"]525            subfolder = img_info.get("subfolder", "")526            img_type = img_info.get("type", "output")527528            resp = requests.get(529                f"{COMFYUI_API_URL}/view",530                params={531                    "filename": filename,532                    "subfolder": subfolder,533                    "type": img_type,534                },535                timeout=30,536            )537            resp.raise_for_status()538539            out_filename = f"edited_{filename}"540            out_path = os.path.join(output_dir, out_filename)541            with open(out_path, "wb") as f:542                f.write(resp.content)543544            return out_filename545546    raise RuntimeError(547        f"No output image found in ComfyUI response. Available outputs: {json.dumps(outputs, indent=2)}"548    )549550551def main():552    try:553        input_json = sys.stdin.read()554        execution_input = json.loads(input_json)555        inputs = execution_input.get("inputs", {})556557        images = inputs.get("images", [])558        prompt = inputs.get("prompt", "")559        aspect_ratio = float(inputs.get("aspect_ratio", 1.0))560561        # Normalize images to a list (single edge gives a string, array edge gives a list)562        if isinstance(images, str):563            images = [images]564565        if not prompt:566            raise ValueError("Prompt is required")567        if not (0.25 <= aspect_ratio <= 4.0):568            raise ValueError(f"Aspect ratio must be between 0.25 and 4.0, got {aspect_ratio}")569        if not images or len(images) == 0:570            raise ValueError("At least one input image is required")571        if len(images) > 2:572            raise ValueError("Maximum of 2 input images supported")573574        os.makedirs(OUTPUT_DIR, exist_ok=True)575576        # Upload input images to ComfyUI577        comfyui_image_names = []578        for img_filename in images:579            local_path = os.path.join(INPUT_DIR, img_filename)580            if not os.path.exists(local_path):581                raise FileNotFoundError(f"Input image not found: {local_path}")582            uploaded_name = upload_image_to_comfyui(local_path)583            comfyui_image_names.append(uploaded_name)584585        # Build workflow based on image count586        if len(comfyui_image_names) == 1:587            workflow = build_single_workflow(prompt, aspect_ratio, comfyui_image_names[0])588        else:589            workflow = build_multiple_workflow(590                prompt, aspect_ratio,591                comfyui_image_names[0], comfyui_image_names[1],592            )593594        # Submit and wait595        prompt_id = submit_prompt(workflow)596        prompt_data = wait_for_result(prompt_id)597        out_filename = download_output_image(prompt_data, OUTPUT_DIR)598599        # Log metadata to stderr600        workflow_type = "single" if len(comfyui_image_names) == 1 else "multiple"601        print(f"prompt_id={prompt_id}, aspect_ratio={aspect_ratio}, images={len(comfyui_image_names)}, workflow={workflow_type}", file=sys.stderr)602603        # Flat output — keys match OUTPUT_SCHEMA604        output = {605            "image": out_filename,606        }607        print(json.dumps(output, indent=2))608609    except Exception as e:610        error_output = {611            "error": str(e),612            "errorType": type(e).__name__,613            "traceback": traceback.format_exc(),614        }615        print(json.dumps(error_output), file=sys.stderr)616        sys.exit(1)617618619if __name__ == "__main__":620    main()