Image Creation 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")18OUTPUT_DIR = "/data/output"1920WORKFLOW = {21    "76": {22        "inputs": {"value": ""},23        "class_type": "PrimitiveStringMultiline",24        "_meta": {"title": "Prompt"},25    },26    "78": {27        "inputs": {"filename_prefix": "Flux2-Klein", "images": ["77:65", 0]},28        "class_type": "SaveImage",29        "_meta": {"title": "Save Image"},30    },31    "84": {32        "inputs": {33            "megapixel": "1.0",34            "aspect_ratio": "9:16 (Slim Vertical)",35            "divisible_by": "64",36            "custom_ratio": True,37            "custom_aspect_ratio": "1:1",38        },39        "class_type": "FluxResolutionNode",40        "_meta": {"title": "Flux Resolution Calc"},41    },42    "77:61": {43        "inputs": {"sampler_name": "euler"},44        "class_type": "KSamplerSelect",45        "_meta": {"title": "KSamplerSelect"},46    },47    "77:64": {48        "inputs": {49            "noise": ["77:73", 0],50            "guider": ["77:63", 0],51            "sampler": ["77:61", 0],52            "sigmas": ["77:62", 0],53            "latent_image": ["77:66", 0],54        },55        "class_type": "SamplerCustomAdvanced",56        "_meta": {"title": "SamplerCustomAdvanced"},57    },58    "77:65": {59        "inputs": {"samples": ["77:64", 0], "vae": ["77:72", 0]},60        "class_type": "VAEDecode",61        "_meta": {"title": "VAE Decode"},62    },63    "77:66": {64        "inputs": {65            "width": ["77:68", 0],66            "height": ["77:69", 0],67            "batch_size": 1,68        },69        "class_type": "EmptyFlux2LatentImage",70        "_meta": {"title": "Empty Flux 2 Latent"},71    },72    "77:68": {73        "inputs": {"value": ["84", 0]},74        "class_type": "PrimitiveInt",75        "_meta": {"title": "Width"},76    },77    "77:69": {78        "inputs": {"value": ["84", 1]},79        "class_type": "PrimitiveInt",80        "_meta": {"title": "Height"},81    },82    "77:73": {83        "inputs": {"noise_seed": 0},84        "class_type": "RandomNoise",85        "_meta": {"title": "RandomNoise"},86    },87    "77:70": {88        "inputs": {89            "unet_name": "flux-2-klein-9b-fp8.safetensors",90            "weight_dtype": "default",91        },92        "class_type": "UNETLoader",93        "_meta": {"title": "Load Diffusion Model"},94    },95    "77:71": {96        "inputs": {97            "clip_name": "qwen_3_8b_fp8mixed.safetensors",98            "type": "flux2",99            "device": "default",100        },101        "class_type": "CLIPLoader",102        "_meta": {"title": "Load CLIP"},103    },104    "77:72": {105        "inputs": {"vae_name": "flux2-vae.safetensors"},106        "class_type": "VAELoader",107        "_meta": {"title": "Load VAE"},108    },109    "77:63": {110        "inputs": {111            "cfg": 1,112            "model": ["77:70", 0],113            "positive": ["77:74", 0],114            "negative": ["77:76", 0],115        },116        "class_type": "CFGGuider",117        "_meta": {"title": "CFGGuider"},118    },119    "77:76": {120        "inputs": {"conditioning": ["77:74", 0]},121        "class_type": "ConditioningZeroOut",122        "_meta": {"title": "ConditioningZeroOut"},123    },124    "77:74": {125        "inputs": {"text": ["76", 0], "clip": ["77:71", 0]},126        "class_type": "CLIPTextEncode",127        "_meta": {"title": "CLIP Text Encode (Positive Prompt)"},128    },129    "77:62": {130        "inputs": {131            "steps": 4,132            "width": ["77:68", 0],133            "height": ["77:69", 0],134        },135        "class_type": "Flux2Scheduler",136        "_meta": {"title": "Flux2Scheduler"},137    },138}139140141def aspect_ratio_to_string(ratio: float) -> str:142    """Convert a float w/h ratio to a W:H string."""143    frac = Fraction(ratio).limit_denominator(64)144    return f"{frac.numerator}:{frac.denominator}"145146147def build_workflow(prompt: str, aspect_ratio: float, megapixel: float = 1.0) -> dict:148    """Inject user inputs into the workflow template."""149    import copy150    wf = copy.deepcopy(WORKFLOW)151152    # Set prompt153    wf["76"]["inputs"]["value"] = prompt154155    # Set aspect ratio156    ratio_str = aspect_ratio_to_string(aspect_ratio)157    wf["84"]["inputs"]["custom_ratio"] = True158    wf["84"]["inputs"]["custom_aspect_ratio"] = ratio_str159160    # Set megapixel161    wf["84"]["inputs"]["megapixel"] = str(megapixel)162163    # Randomize noise seed164    wf["77:73"]["inputs"]["noise_seed"] = random.randint(0, 2**31 - 1)165166    return wf167168169def submit_prompt(workflow: dict) -> str:170    """Submit workflow to ComfyUI and return prompt_id."""171    resp = requests.post(172        f"{COMFYUI_API_URL}/prompt",173        json={"prompt": workflow},174        timeout=30,175    )176    if resp.status_code != 200:177        try:178            error_detail = resp.json()179        except Exception:180            error_detail = resp.text181        raise RuntimeError(182            f"ComfyUI /prompt returned {resp.status_code}: {json.dumps(error_detail, indent=2) if isinstance(error_detail, dict) else error_detail}"183        )184    data = resp.json()185186    # ComfyUI returns 200 even when nodes have validation errors187    node_errors = data.get("node_errors", {})188    if node_errors:189        raise RuntimeError(190            f"ComfyUI workflow has node errors: {json.dumps(node_errors, indent=2)}"191        )192193    return data["prompt_id"]194195196def wait_for_result(prompt_id: str, timeout: int = 600, poll_interval: int = 2) -> dict:197    """Poll ComfyUI history until the prompt completes with outputs."""198    deadline = time.time() + timeout199    empty_complete_retries = 0200    max_empty_retries = 3  # grace period for output serialization lag201202    while time.time() < deadline:203        resp = requests.get(204            f"{COMFYUI_API_URL}/history/{prompt_id}",205            timeout=10,206        )207        resp.raise_for_status()208        history = resp.json()209210        if prompt_id in history:211            prompt_data = history[prompt_id]212            status = prompt_data.get("status", {})213214            if status.get("status_str") == "error":215                messages = status.get("messages", [])216                raise RuntimeError(217                    f"ComfyUI prompt failed: {json.dumps(messages, indent=2)}"218                )219220            if status.get("completed", False):221                if prompt_data.get("outputs"):222                    return prompt_data223224                # Completed but no outputs — retry briefly for race condition225                empty_complete_retries += 1226                if empty_complete_retries >= max_empty_retries:227                    raise RuntimeError(228                        f"ComfyUI prompt completed but produced no outputs. "229                        f"This usually means a node failed silently (missing custom node or model). "230                        f"Status: {json.dumps(status, indent=2)}"231                    )232233        time.sleep(poll_interval)234235    raise TimeoutError(f"ComfyUI prompt {prompt_id} did not complete within {timeout}s")236237238def download_output_image(prompt_data: dict, output_dir: str) -> str:239    """Download the generated image from ComfyUI."""240    outputs = prompt_data.get("outputs", {})241    for node_id, node_output in outputs.items():242        images = node_output.get("images", [])243        if images:244            img_info = images[0]245            filename = img_info["filename"]246            subfolder = img_info.get("subfolder", "")247            img_type = img_info.get("type", "output")248249            resp = requests.get(250                f"{COMFYUI_API_URL}/view",251                params={252                    "filename": filename,253                    "subfolder": subfolder,254                    "type": img_type,255                },256                timeout=30,257            )258            resp.raise_for_status()259260            out_filename = f"generated_{filename}"261            out_path = os.path.join(output_dir, out_filename)262            with open(out_path, "wb") as f:263                f.write(resp.content)264265            return out_filename266267    raise RuntimeError(268        f"No output image found in ComfyUI response. Available outputs: {json.dumps(outputs, indent=2)}"269    )270271272def main():273    try:274        input_json = sys.stdin.read()275        execution_input = json.loads(input_json)276        inputs = execution_input.get("inputs", {})277278        prompt = inputs.get("prompt", "")279        aspect_ratio = float(inputs.get("aspect_ratio", 1.0))280        megapixel = float(inputs.get("megapixel", 1.0))281282        if not prompt:283            raise ValueError("Prompt is required")284        if not (0.25 <= aspect_ratio <= 4.0):285            raise ValueError(f"Aspect ratio must be between 0.25 and 4.0, got {aspect_ratio}")286        if not (0.25 <= megapixel <= 4.0):287            raise ValueError(f"Megapixel must be between 0.25 and 4.0, got {megapixel}")288289        os.makedirs(OUTPUT_DIR, exist_ok=True)290291        # Build and submit workflow292        workflow = build_workflow(prompt, aspect_ratio, megapixel)293        prompt_id = submit_prompt(workflow)294295        # Wait for completion and download result296        prompt_data = wait_for_result(prompt_id)297        out_filename = download_output_image(prompt_data, OUTPUT_DIR)298299        # Log metadata to stderr300        print(f"prompt_id={prompt_id}, aspect_ratio={aspect_ratio}, megapixel={megapixel}", file=sys.stderr)301302        # Flat output — keys match OUTPUT_SCHEMA303        output = {304            "image": out_filename,305        }306        print(json.dumps(output, indent=2))307308    except Exception as e:309        error_output = {310            "error": str(e),311            "errorType": type(e).__name__,312            "traceback": traceback.format_exc(),313        }314        print(json.dumps(error_output), file=sys.stderr)315        sys.exit(1)316317318if __name__ == "__main__":319    main()