from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import requests, base64, time, uuid, random, copy

app = FastAPI()

app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)

COMFY_URL = "http://127.0.0.1:8188"

# 🔥 這份是依你貼的 history「prompt[2]」整理出來（節點ID要完全一致）
ZIMAGE_WORKFLOW = {
"9": {
"inputs": {
"filename_prefix": "z-image-turbo",
"images": ["57:8", 0]
},
"class_type": "SaveImage"
},
"57:30": {
"inputs": {
"clip_name": "qwen_3_4b.safetensors",
"type": "lumina2",
"device": "default"
},
"class_type": "CLIPLoader"
},
"57:29": {
"inputs": {
"vae_name": "ae.safetensors"
},
"class_type": "VAELoader"
},
"57:33": {
"inputs": {
"conditioning": ["57:27", 0]
},
"class_type": "ConditioningZeroOut"
},
"57:8": {
"inputs": {
"samples": ["57:3", 0],
"vae": ["57:29", 0]
},
"class_type": "VAEDecode"
},
"57:28": {
"inputs": {
"unet_name": "z_image_turbo_bf16.safetensors",
"weight_dtype": "default"
},
"class_type": "UNETLoader"
},
"57:27": {
"inputs": {
"text": "test",
"clip": ["57:30", 0]
},
"class_type": "CLIPTextEncode"
},
"57:13": {
"inputs": {
"width": 1024,
"height": 1024,
"batch_size": 1
},
"class_type": "EmptySD3LatentImage"
},
"57:11": {
"inputs": {
"shift": 3.0,
"model": ["57:28", 0]
},
"class_type": "ModelSamplingAuraFlow"
},
"57:3": {
"inputs": {
"seed": 1,
"steps": 8,
"cfg": 1.0,
"sampler_name": "res_multistep",
"scheduler": "simple",
"denoise": 1.0,
"model": ["57:11", 0],
"positive": ["57:27", 0],
"negative": ["57:33", 0],
"latent_image": ["57:13", 0]
},
"class_type": "KSampler"
}
}


@app.post("/generate")
async def generate(data: dict):
user_prompt = (data.get("prompt") or "").strip()
if not user_prompt:
user_prompt = "a cute puppy"

print("\n==== z-image API ====")
print("prompt:", user_prompt)

# 🔥 用你成功的 workflow（深拷貝）
workflow = copy.deepcopy(ZIMAGE_WORKFLOW)

# 🔥 注入 prompt（你的節點ID）
workflow["57:27"]["inputs"]["text"] = user_prompt

# 🔥 隨機 seed（你的節點ID）
workflow["57:3"]["inputs"]["seed"] = random.randint(1, 999999999)

payload = {
"prompt": workflow,
"client_id": str(uuid.uuid4())
}

# 🔥 送出
res = requests.post(f"{COMFY_URL}/prompt", json=payload)
print("RAW:", res.text)

try:
result = res.json()
except:
return {"error": "ComfyUI 回傳非JSON"}

prompt_id = result.get("prompt_id")
if not prompt_id:
return {"error": result} # 直接把錯誤回傳出來

# 🔥 等結果（用 history 查）
for i in range(30):
time.sleep(1)

hist = requests.get(f"{COMFY_URL}/history/{prompt_id}").json()

if prompt_id in hist:
outputs = hist[prompt_id].get("outputs", {})

for node_id in outputs:
images = outputs[node_id].get("images", [])
if images:
img = images[0]
filename = img["filename"]
subfolder = img.get("subfolder", "")
img_type = img.get("type", "output")

url = f"{COMFY_URL}/view?filename={filename}&subfolder={subfolder}&type={img_type}"
print("✅ URL:", url)

img_bytes = requests.get(url).content
img_base64 = base64.b64encode(img_bytes).decode("utf-8")

return {"image_base64": img_base64}

print(f"等待中... {i+1}s")

return {"error": "沒有生成圖片"}

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import requests, base64, os, time, uuid

app = FastAPI()

app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)

COMFY_URL = "http://127.0.0.1:8188"
OUTPUT_DIR = "/home/jetbomb/ComfyUI/output"


# 🔥 抓最新 workflow（關鍵）
def get_latest_workflow():
res = requests.get(f"{COMFY_URL}/history")
history = res.json()

# 找最新一筆
latest = list(history.values())[-1]

# 取 prompt（真正可用）
workflow = latest["prompt"][2]

return workflow


@app.post("/generate")
async def generate(data: dict):

user_prompt = data.get("prompt", "").strip()
if not user_prompt:
user_prompt = "product"

print("\n==== 新請求 ====")
print("輸入:", user_prompt)

# 🔥 取得你UI正在用的 workflow
workflow = get_latest_workflow()

# 🔥 找 prompt node（通常是89）
for node_id, node in workflow.items():
if node.get("class_type") == "CLIPTextEncode":
if "text" in node["inputs"]:
node["inputs"]["text"] = f"""
{user_prompt},

product photography,
clean white background,
studio lighting,
realistic,
high detail,
no logo, no text
"""

# 只改第一個（正向）
break

# 🔥 記錄舊檔案
before = set(os.listdir(OUTPUT_DIR))

payload = {
"prompt": workflow,
"client_id": str(uuid.uuid4())
}

res = requests.post(f"{COMFY_URL}/prompt", json=payload)
print("ComfyUI 回應:", res.text)

# 🔥 等新圖
for i in range(40):
time.sleep(1)

after = set(os.listdir(OUTPUT_DIR))
new_files = after - before

pngs = [f for f in new_files if f.endswith(".png")]

if pngs:
latest = sorted(
pngs,
key=lambda x: os.path.getmtime(os.path.join(OUTPUT_DIR, x)),
reverse=True
)[0]

path = os.path.join(OUTPUT_DIR, latest)

print("✅ 抓到:", path)

with open(path, "rb") as f:
img_base64 = base64.b64encode(f.read()).decode("utf-8")

return {"image_base64": img_base64}

print(f"等待中... {i+1}s")

return {"error": "沒有生成圖片"}