Google AI studio + FASTAPI + ComfyUI 연동

import base64
import io
import json
import time
from pathlib import Path
from typing import Optional, Dict, Any

import httpx
from fastapi import FastAPI, HTTPException, Response, Header
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from PIL import Image

COMFY_URL = "http://127.0.0.1:8000"
WORKFLOW_PATH = Path("./img2img.json")

API_KEY = "a1s2d3f4"  # 실험용. 나중에 바꾸면 됨.

MAX_IMAGE_BYTES = 16 * 1024 * 1024  # input base64 업로드 제한(16MB)
DEFAULT_JPEG_QUALITY = 85
DEFAULT_MAX_EDGE = 1024  # JPEG로 받을 때 큰 이미지면 리사이즈(전송/속도 개선)

app = FastAPI(title="ComfyUI img2img API (async + JPEG result)")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)


class Img2ImgRequest(BaseModel):
    # 입력(둘 중 하나)
    input_filename: Optional[str] = None
    input_image_base64: Optional[str] = None

    # 프롬프트 수정
    positive_text: Optional[str] = None
    positive_add: str = ""
    negative_text: Optional[str] = None
    negative_add: str = ""

    # 샘플링 파라미터
    steps: int = Field(30, ge=1, le=80)
    cfg: float = Field(6.5, ge=1.0, le=15.0)
    start_at_step: int = Field(16, ge=0, le=79)
    end_at_step: int = Field(30, ge=1, le=80)
    seed: int = Field(-1, description="-1=random, >=0 fixed")

    filename_prefix: str = "ComfyUI_API"


def _auth(x_api_key: Optional[str]) -> None:
    if x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")


def _load_workflow() -> Dict[str, Any]:
    if not WORKFLOW_PATH.exists():
        raise HTTPException(
            status_code=500,
            detail=f"Workflow JSON not found: {WORKFLOW_PATH.resolve()}",
        )
    return json.loads(WORKFLOW_PATH.read_text(encoding="utf-8"))


async def _upload_to_comfy_input(client: httpx.AsyncClient, img_b64: str, filename: str = "input.png") -> str:
    try:
        img_bytes = base64.b64decode(img_b64)
    except Exception:
        raise HTTPException(status_code=400, detail="input_image_base64 is not valid base64")

    if len(img_bytes) > MAX_IMAGE_BYTES:
        raise HTTPException(status_code=413, detail=f"Image too large (>{MAX_IMAGE_BYTES} bytes).")

    files = {"image": (filename, img_bytes, "application/octet-stream")}
    data = {"type": "input"}  # ComfyUI input 폴더에 업로드
    r = await client.post(f"{COMFY_URL}/upload/image", data=data, files=files, timeout=120)
    r.raise_for_status()
    resp = r.json()
    return resp.get("name", filename)


def _patch_workflow(wf: Dict[str, Any], req: Img2ImgRequest, input_filename: str) -> Dict[str, Any]:
    """
    이 patch는 네가 준 워크플로우 JSON(node id)에 맞춰서 수정됨.
    필요하면 node id만 바꿔주면 됨.
    """
    NODE_SAVE = "6"
    NODE_LOAD = "11"
    NODE_POS = "13"
    NODE_NEG = "14"
    NODE_SAMPLER = "10"

    # LoadImage 노드에 입력 파일명 설정
    wf[NODE_LOAD]["inputs"]["image"] = input_filename

    # Positive prompt
    if req.positive_text is not None:
        wf[NODE_POS]["inputs"]["text"] = req.positive_text
    elif req.positive_add.strip():
        wf[NODE_POS]["inputs"]["text"] = wf[NODE_POS]["inputs"]["text"].rstrip() + "\n" + req.positive_add.strip() + "\n"

    # Negative prompt
    if req.negative_text is not None:
        wf[NODE_NEG]["inputs"]["text"] = req.negative_text
    elif req.negative_add.strip():
        wf[NODE_NEG]["inputs"]["text"] = wf[NODE_NEG]["inputs"]["text"].rstrip() + "\n" + req.negative_add.strip() + "\n"

    # Sampler params
    s = wf[NODE_SAMPLER]["inputs"]
    s["steps"] = int(req.steps)
    s["cfg"] = float(req.cfg)
    s["start_at_step"] = int(req.start_at_step)
    s["end_at_step"] = int(req.end_at_step)

    # seed
    if req.seed < 0:
        s["add_noise"] = "enable"
        s["noise_seed"] = int(time.time() * 1000) % 2_147_483_647
    else:
        s["add_noise"] = "enable"
        s["noise_seed"] = int(req.seed)

    if "return_with_leftover_noise" in s:
        s["return_with_leftover_noise"] = "disable"

    # filename prefix
    wf[NODE_SAVE]["inputs"]["filename_prefix"] = req.filename_prefix
    return wf


async def _comfy_enqueue(client: httpx.AsyncClient, wf: Dict[str, Any]) -> str:
    r = await client.post(f"{COMFY_URL}/prompt", json={"prompt": wf}, timeout=120)
    r.raise_for_status()
    return r.json()["prompt_id"]


async def _get_history_item(client: httpx.AsyncClient, prompt_id: str) -> Optional[Dict[str, Any]]:
    r = await client.get(f"{COMFY_URL}/history/{prompt_id}", timeout=30)
    r.raise_for_status()
    hist = r.json()
    return hist.get(prompt_id)


def _extract_first_image_meta(history_item: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    outputs = history_item.get("outputs", {})
    for _, out in outputs.items():
        images = out.get("images")
        if images:
            return images[0]
    return None


async def _fetch_image_bytes(client: httpx.AsyncClient, meta: Dict[str, Any]) -> bytes:
    filename = meta.get("filename")
    subfolder = meta.get("subfolder", "")
    img_type = meta.get("type", "output")
    if not filename:
        raise HTTPException(status_code=500, detail="No filename in ComfyUI image metadata")

    vr = await client.get(
        f"{COMFY_URL}/view",
        params={"filename": filename, "subfolder": subfolder, "type": img_type},
        timeout=120,
    )
    vr.raise_for_status()
    return vr.content


def _png_bytes_to_jpeg_bytes(png_bytes: bytes, quality: int = DEFAULT_JPEG_QUALITY, max_edge: int = DEFAULT_MAX_EDGE) -> bytes:
    img = Image.open(io.BytesIO(png_bytes)).convert("RGB")

    if max_edge and max(img.size) > max_edge:
        w, h = img.size
        if w >= h:
            new_w = max_edge
            new_h = int(h * (max_edge / w))
        else:
            new_h = max_edge
            new_w = int(w * (max_edge / h))
        img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)

    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=int(quality), optimize=True)
    return buf.getvalue()


@app.get("/")
def root():
    return {"ok": True, "endpoints": ["/docs", "POST /img2img", "GET /result/{prompt_id}"]}


@app.post("/img2img")
async def img2img_enqueue(req: Img2ImgRequest, x_api_key: Optional[str] = Header(None)):
    """
    비동기 enqueue: 바로 prompt_id만 반환 (Cloudflare 524 회피)
    """
    _auth(x_api_key)
    wf = _load_workflow()

    async with httpx.AsyncClient() as client:
        # 입력 결정
        if req.input_image_base64:
            input_name = await _upload_to_comfy_input(client, req.input_image_base64, filename="input.png")
        elif req.input_filename:
            input_name = req.input_filename
        else:
            raise HTTPException(status_code=400, detail="Provide either input_filename or input_image_base64")

        wf = _patch_workflow(wf, req, input_name)
        prompt_id = await _comfy_enqueue(client, wf)

    return {"prompt_id": prompt_id, "status": "queued"}


@app.get("/result/{prompt_id}")
async def result(
    prompt_id: str,
    x_api_key: Optional[str] = Header(None),
    base64_out: bool = False,
    jpeg: bool = True,
    quality: int = DEFAULT_JPEG_QUALITY,
    max_edge: int = DEFAULT_MAX_EDGE,
):
    """
    폴링:
    - running이면 {"status":"running"}
    - done이면 기본은 JPEG 바이너리(image/jpeg)로 반환
    - base64_out=true면 JSON으로 base64 반환 (jpeg=true면 JPEG base64)
    """
    _auth(x_api_key)

    async with httpx.AsyncClient() as client:
        item = await _get_history_item(client, prompt_id)
        if not item:
            return {"prompt_id": prompt_id, "status": "running"}

        meta = _extract_first_image_meta(item)
        if not meta:
            return {"prompt_id": prompt_id, "status": "running"}

        png_bytes = await _fetch_image_bytes(client, meta)

    out_bytes = png_bytes
    out_media = "image/png"
    out_format = "png"

    if jpeg:
        out_bytes = _png_bytes_to_jpeg_bytes(png_bytes, quality=quality, max_edge=max_edge)
        out_media = "image/jpeg"
        out_format = "jpeg"

    if base64_out:
        return {
            "prompt_id": prompt_id,
            "status": "done",
            "format": out_format,
            "image_base64": base64.b64encode(out_bytes).decode("utf-8"),
        }

    return Response(content=out_bytes, media_type=out_media)

cloudflared tunnel --url <http://127.0.0.1:9001>

You are building a lightweight image transformation client app.

This app does NOT generate images by itself.
Instead, it acts as a frontend client that communicates with an external image processing API.

The external backend is:
- A FastAPI server
- Running locally with ComfyUI
- Exposed via a Cloudflare Tunnel
- The base URL of the backend CHANGES every time the tunnel restarts

Your job is to build a frontend app with the following requirements.

---

[1] Settings Panel (Important)

The app must provide a settings section where the user can manually input:

- Backend Base URL (example: https://xxxx.trycloudflare.com)
- API Key (string)

These values must:
- Be editable by the user
- Be stored in local state (or localStorage)
- Be used dynamically for all API calls
- NOT be hardcoded

Include a simple “Test Connection” button that sends a GET request to:
{BaseURL}/docs
and shows success or failure.

---

[2] Main Workflow: Image-to-Image

The app workflow is:

1. User uploads an image (JPEG or PNG)
2. User optionally enters a text prompt describing how the image should be modified
3. When the user clicks “Transform”:
   - Send a POST request to:
     {BaseURL}/img2img
   - Include:
     - the uploaded image
     - prompt text
     - any parameters as JSON
   - Include header:
     X-API-Key: {API Key}

4. The server returns:
   {
     "prompt_id": "string",
     "status": "queued"
   }

5. The app must then poll:
   GET {BaseURL}/result/{prompt_id}?base64_out=true&jpeg=true

6. Poll until status == "done"

7. When done:
   - Receive a JSON response with:
     {
       "format": "jpeg",
       "image_base64": "..."
     }
   - Decode and display the resulting image

---

[3] UI Expectations

- Simple, minimal UI
- Left or top area: settings (Base URL, API Key)
- Main area:
  - Image upload
  - Prompt input
  - Transform button
  - Result image preview
- Show basic status:
  - queued
  - processing
  - done
  - error / timeout

---

[4] Important Constraints

- Do NOT assume the backend URL is stable
- Do NOT embed backend logic
- Do NOT generate images locally
- The app is only a client

The purpose of this app is to allow rapid experimentation with ComfyUI image-to-image workflows running on a local machine, accessed remotely via Cloudflare Tunnel.

Build the app accordingly.