Image Remove Background Legacy

1import os2import sys3import json4import time5import traceback67try:8    import requests9except ImportError:10    import subprocess11    subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])12    import requests1314COMFYUI_API_URL = os.getenv("COMFYUI_API_URL", "http://192.168.1.39:8188")15INPUT_DIR = "/data/input"16OUTPUT_DIR = "/data/output"1718WORKFLOW = {19    "1": {20        "inputs": {21            "model": "RMBG-2.0",22            "sensitivity": 1,23            "process_res": 1024,24            "mask_blur": 0,25            "mask_offset": 0,26            "invert_output": True,27            "refine_foreground": False,28            "background": "Alpha",29            "background_color": "#222222",30            "image": ["7", 0],31        },32        "class_type": "RMBG",33        "_meta": {"title": "Remove Background (RMBG)"},34    },35    "6": {36        "inputs": {37            "filename_prefix": "removebg_emblema",38            "images": ["1", 0],39            "mask": ["1", 1],40        },41        "class_type": "SaveImageWithAlpha",42        "_meta": {"title": "Save Image With Alpha"},43    },44    "7": {45        "inputs": {"image": ""},46        "class_type": "LoadImage",47        "_meta": {"title": "Load Image"},48    },49}505152def upload_file_to_comfyui(local_path: str) -> str:53    """Upload a local image to ComfyUI and return the uploaded filename."""54    with open(local_path, "rb") as f:55        resp = requests.post(56            f"{COMFYUI_API_URL}/upload/image",57            files={"image": (os.path.basename(local_path), f, "image/png")},58            timeout=30,59        )60    resp.raise_for_status()61    data = resp.json()62    return data["name"]636465def build_workflow(image_name: str) -> dict:66    """Inject the uploaded image name into the workflow template."""67    import copy68    wf = copy.deepcopy(WORKFLOW)69    wf["7"]["inputs"]["image"] = image_name70    wf["6"]["inputs"]["filename_prefix"] = "removebg_emblema"71    return wf727374def submit_prompt(workflow: dict) -> str:75    """Submit workflow to ComfyUI and return prompt_id."""76    resp = requests.post(77        f"{COMFYUI_API_URL}/prompt",78        json={"prompt": workflow},79        timeout=30,80    )81    if resp.status_code != 200:82        try:83            error_detail = resp.json()84        except Exception:85            error_detail = resp.text86        raise RuntimeError(87            f"ComfyUI /prompt returned {resp.status_code}: {json.dumps(error_detail, indent=2) if isinstance(error_detail, dict) else error_detail}"88        )89    data = resp.json()9091    node_errors = data.get("node_errors", {})92    if node_errors:93        raise RuntimeError(94            f"ComfyUI workflow has node errors: {json.dumps(node_errors, indent=2)}"95        )9697    return data["prompt_id"]9899100def wait_for_result(prompt_id: str, timeout: int = 600, poll_interval: int = 2) -> dict:101    """Poll ComfyUI history until the prompt completes with outputs."""102    deadline = time.time() + timeout103    empty_complete_retries = 0104    max_empty_retries = 3105106    while time.time() < deadline:107        resp = requests.get(108            f"{COMFYUI_API_URL}/history/{prompt_id}",109            timeout=10,110        )111        resp.raise_for_status()112        history = resp.json()113114        if prompt_id in history:115            prompt_data = history[prompt_id]116            status = prompt_data.get("status", {})117118            if status.get("status_str") == "error":119                messages = status.get("messages", [])120                raise RuntimeError(121                    f"ComfyUI prompt failed: {json.dumps(messages, indent=2)}"122                )123124            if status.get("completed", False):125                if prompt_data.get("outputs"):126                    return prompt_data127128                empty_complete_retries += 1129                if empty_complete_retries >= max_empty_retries:130                    raise RuntimeError(131                        f"ComfyUI prompt completed but produced no outputs. "132                        f"This usually means a node failed silently (missing custom node or model). "133                        f"Status: {json.dumps(status, indent=2)}"134                    )135136        time.sleep(poll_interval)137138    raise TimeoutError(f"ComfyUI prompt {prompt_id} did not complete within {timeout}s")139140141def download_output_image(prompt_data: dict, output_dir: str) -> str:142    """Download the output PNG from ComfyUI."""143    outputs = prompt_data.get("outputs", {})144    for node_id, node_output in outputs.items():145        images = node_output.get("images", [])146        if images:147            img_info = images[0]148            filename = img_info["filename"]149            subfolder = img_info.get("subfolder", "")150            img_type = img_info.get("type", "output")151152            resp = requests.get(153                f"{COMFYUI_API_URL}/view",154                params={155                    "filename": filename,156                    "subfolder": subfolder,157                    "type": img_type,158                },159                timeout=30,160            )161            resp.raise_for_status()162163            out_filename = f"removebg_{filename}"164            out_path = os.path.join(output_dir, out_filename)165            with open(out_path, "wb") as f:166                f.write(resp.content)167168            return out_filename169170    raise RuntimeError(171        f"No output image found in ComfyUI response. Available outputs: {json.dumps(outputs, indent=2)}"172    )173174175def main():176    try:177        input_json = sys.stdin.read()178        execution_input = json.loads(input_json)179        inputs = execution_input.get("inputs", {})180181        image = inputs.get("image", "")182        if not image:183            raise ValueError("Input image is required")184185        os.makedirs(OUTPUT_DIR, exist_ok=True)186187        # Upload input image to ComfyUI188        local_path = os.path.join(INPUT_DIR, image)189        if not os.path.exists(local_path):190            raise FileNotFoundError(f"Input image not found: {local_path}")191        uploaded_name = upload_file_to_comfyui(local_path)192193        # Build and submit workflow194        workflow = build_workflow(uploaded_name)195        prompt_id = submit_prompt(workflow)196197        # Wait for completion and download result198        prompt_data = wait_for_result(prompt_id)199        out_filename = download_output_image(prompt_data, OUTPUT_DIR)200201        # Log metadata to stderr202        print(f"prompt_id={prompt_id}", file=sys.stderr)203204        # Flat output — keys match OUTPUT_SCHEMA205        output = {206            "image": out_filename,207        }208        print(json.dumps(output, indent=2))209210    except Exception as e:211        error_output = {212            "error": str(e),213            "errorType": type(e).__name__,214            "traceback": traceback.format_exc(),215        }216        print(json.dumps(error_output), file=sys.stderr)217        sys.exit(1)218219220if __name__ == "__main__":221    main()