cbpark84-ComfyBridgeIO/examples/batch_example.py

"""배치 워크플로우 실행 예제."""

import asyncio
import json
from pathlib import Path

from comfybridgeio import ComfyBridge, ResultCache


async def batch_sequential_example():
    """
    여러 프롬프트를 순차적으로 처리한다.
    비디오 출력이 있으면 자동으로 순차 모드 사용.
    """
    bridge = ComfyBridge(host="localhost", port=8188)
    cache = ResultCache(max_size=50, ttl=1800)

    prompts = [
        {"prompt": "a red cat", "steps": 20},
        {"prompt": "a blue dog", "steps": 20},
        {"prompt": "a green bird", "steps": 20},
    ]

    results = []
    for i, prompt_dict in enumerate(prompts):
        print(f"[{i+1}/{len(prompts)}] Running {prompt_dict['prompt']}")

        result = await bridge.run_batch(
            workflow="workflow.json",
            inputs=[prompt_dict],
            mode="sequential",
        )
        results.extend(result)

        # 결과 캐시
        cache.set(f"prompt_{i}", result[0] if result else None)

    print(f"\nGenerated {len(results)} images")
    for result in results:
        print(f"  - {result.keys()}")

    return results


async def batch_concurrent_example():
    """
    여러 프롬프트를 동시에 처리한다.
    텍스트 생성 등 빠른 작업에 권장.
    """
    bridge = ComfyBridge(host="localhost", port=8188)

    prompts = [
        {"negative_prompt": "blurry", "seed": 1},
        {"negative_prompt": "blurry", "seed": 2},
        {"negative_prompt": "blurry", "seed": 3},
        {"negative_prompt": "blurry", "seed": 4},
    ]

    results = await bridge.run_batch(
        workflow="workflow.json",
        inputs=prompts,
        mode="concurrent",
        max_concurrent=4,
    )

    print(f"Generated {len(results)} results in parallel")
    return results


async def chain_workflows_example():
    """워크플로우를 체이닝한다."""
    bridge = ComfyBridge(host="localhost", port=8188)

    # 워크플로우 1: 텍스트 → 이미지
    result1 = await bridge.run(
        "t2i.json", inputs={"prompt": "a cat"}
    )

    # 워크플로우 2: 이미지 → 스타일 전환
    image = result1.as_bytes("output_image")
    uploaded = await bridge.upload_image(image, "style_input.png")

    result2 = await bridge.chain(
        [
            ("t2i.json", {"prompt": "a dog"}),
            ("style.json", {"image": uploaded, "style": "oil_painting"}),
        ],
        progress_callback=lambda p: print(f"Progress: {p:.0%}"),
    )

    print(f"Chain completed: {result2.keys()}")
    return result2


async def upload_files_example():
    """다양한 파일을 업로드한다."""
    bridge = ComfyBridge(host="localhost", port=8188)

    # 이미지
    image_filename = await bridge.upload_image("input.png")
    print(f"Uploaded image: {image_filename}")

    # 오디오
    audio_filename = await bridge.upload_audio("speech.wav")
    print(f"Uploaded audio: {audio_filename}")

    # 비디오
    video_filename = await bridge.upload_video("source.mp4")
    print(f"Uploaded video: {video_filename}")


async def preview_callback_example():
    """
    KSampler 스텝별 중간 이미지를 실시간으로 수신한다.
    """
    bridge = ComfyBridge(host="localhost", port=8188)

    def on_preview(image_bytes: bytes, step: int, total_steps: int):
        print(f"Preview: step {step}/{total_steps}, size={len(image_bytes)} bytes")
        # 이미지 저장, 화면 갱신 등

    result = await bridge.run(
        "workflow.json",
        inputs={"prompt": "a cat", "steps": 20},
        preview_callback=on_preview,
    )

    print(f"Final result: {result.keys()}")


async def extract_schema_example():
    """워크플로우 스키마를 추출하고 검증한다."""
    from comfybridgeio.client.schema import extract_schema, describe_schema, validate_inputs

    workflow = json.load(open("workflow.json"))
    schema = extract_schema(workflow)

    print(describe_schema(schema))

    # 입력 검증
    test_inputs = {"prompt": "test", "steps": 20}
    is_valid, errors = validate_inputs(schema, test_inputs)

    if is_valid:
        print("✓ Inputs are valid")
    else:
        print("✗ Input errors:")
        for err in errors:
            print(f"  - {err}")


if __name__ == "__main__":
    # 예제 선택
    example = "batch_sequential"

    if example == "batch_sequential":
        asyncio.run(batch_sequential_example())
    elif example == "batch_concurrent":
        asyncio.run(batch_concurrent_example())
    elif example == "chain":
        asyncio.run(chain_workflows_example())
    elif example == "upload":
        asyncio.run(upload_files_example())
    elif example == "preview":
        asyncio.run(preview_callback_example())
    elif example == "schema":
        asyncio.run(extract_schema_example())