Image Upscale Legacy

1import os2import sys3import json4import time5import random6import traceback78try:9    import requests10except ImportError:11    import subprocess12    subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])13    import requests1415COMFYUI_API_URL = os.getenv("COMFYUI_API_URL", "http://192.168.1.39:8188")16INPUT_DIR = "/data/input"17OUTPUT_DIR = "/data/output"1819WORKFLOW = {20    "10": {21        "inputs": {"vae_name": "ae.safetensors"},22        "class_type": "VAELoader",23        "_meta": {"title": "Load VAE"},24    },25    "11": {26        "inputs": {27            "clip_name1": "ViT-L-14-TEXT-detail-improved-hiT-GmP-TE-only-HF.safetensors",28            "clip_name2": "t5xxl_fp16.safetensors",29            "type": "flux",30            "device": "default",31        },32        "class_type": "DualCLIPLoader",33        "_meta": {"title": "DualCLIPLoader"},34    },35    "12": {36        "inputs": {37            "unet_name": "flux1-dev.safetensors",38            "weight_dtype": "default",39        },40        "class_type": "UNETLoader",41        "_meta": {"title": "Load Diffusion Model"},42    },43    "22": {44        "inputs": {45            "model": ["12", 0],46            "conditioning": ["40", 0],47        },48        "class_type": "BasicGuider",49        "_meta": {"title": "BasicGuider"},50    },51    "26": {52        "inputs": {"model_name": "4x_NMKD-Siax_200k.pth"},53        "class_type": "UpscaleModelLoader",54        "_meta": {"title": "Load Upscale Model"},55    },56    "29": {57        "inputs": {58            "filename_prefix": "image/emblema-upscale",59            "images": ["39", 0],60        },61        "class_type": "SaveImage",62        "_meta": {"title": "Save Image"},63    },64    "37": {65        "inputs": {66            "text": "blurry",67            "clip": ["11", 0],68        },69        "class_type": "CLIPTextEncode",70        "_meta": {"title": "CLIP Text Encode (Prompt)"},71    },72    "39": {73        "inputs": {74            "upscale_by": ["57", 0],75            "seed": 0,76            "steps": 1,77            "cfg": 2,78            "sampler_name": "euler",79            "scheduler": "normal",80            "denoise": 0.2,81            "mode_type": "Linear",82            "tile_width": ["52", 0],83            "tile_height": ["53", 0],84            "mask_blur": 8,85            "tile_padding": 32,86            "seam_fix_mode": "None",87            "seam_fix_denoise": 1,88            "seam_fix_width": 64,89            "seam_fix_mask_blur": 8,90            "seam_fix_padding": 16,91            "force_uniform_tiles": True,92            "tiled_decode": False,93            "batch_size": 1,94            "image": ["61", 0],95            "model": ["12", 0],96            "positive": ["40", 0],97            "negative": ["37", 0],98            "vae": ["10", 0],99            "upscale_model": ["26", 0],100        },101        "class_type": "UltimateSDUpscale",102        "_meta": {"title": "Ultimate SD Upscale"},103    },104    "40": {105        "inputs": {106            "clip_l": "",107            "t5xxl": "",108            "guidance": 3.5,109            "clip": ["11", 0],110        },111        "class_type": "CLIPTextEncodeFlux",112        "_meta": {"title": "CLIPTextEncodeFlux"},113    },114    "47": {115        "inputs": {"image": ["61", 0]},116        "class_type": "GetImageSize+",117        "_meta": {"title": "Get Image Size"},118    },119    "48": {120        "inputs": {121            "Value_A": ["49", 0],122            "Value_B": ["57", 0],123        },124        "class_type": "DF_Multiply",125        "_meta": {"title": "Multiply"},126    },127    "49": {128        "inputs": {129            "output_type": "float",130            "*": ["47", 0],131        },132        "class_type": "easy convertAnything",133        "_meta": {"title": "Convert Any"},134    },135    "50": {136        "inputs": {137            "output_type": "float",138            "*": ["47", 1],139        },140        "class_type": "easy convertAnything",141        "_meta": {"title": "Convert Any"},142    },143    "51": {144        "inputs": {145            "Value_A": ["50", 0],146            "Value_B": ["57", 0],147        },148        "class_type": "DF_Multiply",149        "_meta": {"title": "Multiply"},150    },151    "52": {152        "inputs": {153            "output_type": "int",154            "*": ["48", 0],155        },156        "class_type": "easy convertAnything",157        "_meta": {"title": "Convert Any"},158    },159    "53": {160        "inputs": {161            "output_type": "int",162            "*": ["51", 0],163        },164        "class_type": "easy convertAnything",165        "_meta": {"title": "Convert Any"},166    },167    "54": {168        "inputs": {169            "number_type": "integer",170            "number": 3.0,171        },172        "class_type": "Constant Number",173        "_meta": {"title": "Constant Number"},174    },175    "57": {176        "inputs": {177            "output_type": "float",178            "*": ["54", 0],179        },180        "class_type": "easy convertAnything",181        "_meta": {"title": "Convert Any"},182    },183    "61": {184        "inputs": {185            "image": "",186            "custom_width": 0,187            "custom_height": 0,188        },189        "class_type": "VHS_LoadImagePath",190        "_meta": {"title": "Load Image (Path)"},191    },192}193194195def upload_image_to_comfyui(local_path: str) -> str:196    """Upload a local image to ComfyUI and return the uploaded filename."""197    with open(local_path, "rb") as f:198        resp = requests.post(199            f"{COMFYUI_API_URL}/upload/image",200            files={"image": (os.path.basename(local_path), f, "image/png")},201            timeout=30,202        )203    resp.raise_for_status()204    data = resp.json()205    return data["name"]206207208def build_workflow(image_name: str, upscale_factor: float) -> dict:209    """Inject user inputs into the workflow template."""210    import copy211    wf = copy.deepcopy(WORKFLOW)212213    # Set input image — swap VHS_LoadImagePath to LoadImage214    wf["61"] = {215        "inputs": {"image": image_name},216        "class_type": "LoadImage",217        "_meta": {"title": "Load Image"},218    }219220    # Set upscale factor221    wf["54"]["inputs"]["number"] = float(upscale_factor)222223    # Randomize seed224    wf["39"]["inputs"]["seed"] = random.randint(0, 2**31 - 1)225226    # Set output prefix227    wf["29"]["inputs"]["filename_prefix"] = "image/emblema-upscale"228229    return wf230231232def submit_prompt(workflow: dict) -> str:233    """Submit workflow to ComfyUI and return prompt_id."""234    resp = requests.post(235        f"{COMFYUI_API_URL}/prompt",236        json={"prompt": workflow},237        timeout=30,238    )239    if resp.status_code != 200:240        try:241            error_detail = resp.json()242        except Exception:243            error_detail = resp.text244        raise RuntimeError(245            f"ComfyUI /prompt returned {resp.status_code}: {json.dumps(error_detail, indent=2) if isinstance(error_detail, dict) else error_detail}"246        )247    data = resp.json()248249    # ComfyUI returns 200 even when nodes have validation errors250    node_errors = data.get("node_errors", {})251    if node_errors:252        raise RuntimeError(253            f"ComfyUI workflow has node errors: {json.dumps(node_errors, indent=2)}"254        )255256    return data["prompt_id"]257258259def wait_for_result(prompt_id: str, timeout: int = 600, poll_interval: int = 2) -> dict:260    """Poll ComfyUI history until the prompt completes with outputs."""261    deadline = time.time() + timeout262    empty_complete_retries = 0263    max_empty_retries = 3  # grace period for output serialization lag264265    while time.time() < deadline:266        resp = requests.get(267            f"{COMFYUI_API_URL}/history/{prompt_id}",268            timeout=10,269        )270        resp.raise_for_status()271        history = resp.json()272273        if prompt_id in history:274            prompt_data = history[prompt_id]275            status = prompt_data.get("status", {})276277            if status.get("status_str") == "error":278                messages = status.get("messages", [])279                raise RuntimeError(280                    f"ComfyUI prompt failed: {json.dumps(messages, indent=2)}"281                )282283            if status.get("completed", False):284                if prompt_data.get("outputs"):285                    return prompt_data286287                # Completed but no outputs — retry briefly for race condition288                empty_complete_retries += 1289                if empty_complete_retries >= max_empty_retries:290                    raise RuntimeError(291                        f"ComfyUI prompt completed but produced no outputs. "292                        f"This usually means a node failed silently (missing custom node or model). "293                        f"Status: {json.dumps(status, indent=2)}"294                    )295296        time.sleep(poll_interval)297298    raise TimeoutError(f"ComfyUI prompt {prompt_id} did not complete within {timeout}s")299300301def download_output_image(prompt_data: dict, output_dir: str) -> str:302    """Download the generated image from ComfyUI."""303    outputs = prompt_data.get("outputs", {})304    for node_id, node_output in outputs.items():305        images = node_output.get("images", [])306        if images:307            img_info = images[0]308            filename = img_info["filename"]309            subfolder = img_info.get("subfolder", "")310            img_type = img_info.get("type", "output")311312            resp = requests.get(313                f"{COMFYUI_API_URL}/view",314                params={315                    "filename": filename,316                    "subfolder": subfolder,317                    "type": img_type,318                },319                timeout=30,320            )321            resp.raise_for_status()322323            out_filename = f"upscaled_{filename}"324            out_path = os.path.join(output_dir, out_filename)325            with open(out_path, "wb") as f:326                f.write(resp.content)327328            return out_filename329330    raise RuntimeError(331        f"No output image found in ComfyUI response. Available outputs: {json.dumps(outputs, indent=2)}"332    )333334335def main():336    try:337        input_json = sys.stdin.read()338        execution_input = json.loads(input_json)339        inputs = execution_input.get("inputs", {})340341        image = inputs.get("image", "")342        upscale_factor = float(inputs.get("upscale_factor", 3))343344        if not image:345            raise ValueError("Image is required")346        if not (2 <= upscale_factor <= 4):347            raise ValueError(f"Upscale factor must be between 2 and 4, got {upscale_factor}")348349        os.makedirs(OUTPUT_DIR, exist_ok=True)350351        # Upload input image to ComfyUI352        local_path = os.path.join(INPUT_DIR, image)353        if not os.path.exists(local_path):354            raise FileNotFoundError(f"Input image not found: {local_path}")355        image_name = upload_image_to_comfyui(local_path)356357        # Build and submit workflow358        workflow = build_workflow(image_name, upscale_factor)359        prompt_id = submit_prompt(workflow)360361        # Wait for completion and download result362        prompt_data = wait_for_result(prompt_id)363        out_filename = download_output_image(prompt_data, OUTPUT_DIR)364365        # Log metadata to stderr366        print(f"prompt_id={prompt_id}, upscale_factor={upscale_factor}", file=sys.stderr)367368        # Flat output — keys match OUTPUT_SCHEMA369        output = {370            "image": out_filename,371        }372        print(json.dumps(output, indent=2))373374    except Exception as e:375        error_output = {376            "error": str(e),377            "errorType": type(e).__name__,378            "traceback": traceback.format_exc(),379        }380        print(json.dumps(error_output), file=sys.stderr)381        sys.exit(1)382383384if __name__ == "__main__":385    main()