"""
mj_imagine → mj_modal 完整链路 demo（局部重绘）
依赖：pip install requests pillow
"""

import base64
import time
from io import BytesIO
import requests
from PIL import Image, ImageDraw
from pprint import pprint

API_KEY = "YOUR API KEY" # ←←← 换成你的
BASE_URL = "https://api.wenwen-ai.com/mj"

---------- 工具函数 ----------

session = requests.Session()
session.headers.update({"Authorization": f"Bearer {API_KEY}"})

def _post(path, json):
resp = session.post(f"{BASE_URL}{path}", json=json)
resp.raise_for_status()
return resp.json()

def _get(path):
resp = session.get(f"{BASE_URL}{path}")
resp.raise_for_status()
return resp.json()

def wait_task(task_id: str):
"""轮询直到任务结束，返回最终 payload"""
while True:
data = _get(f"/task/{task_id}/fetch")
status = data["status"]
print(f"[{task_id}] status = {status}")
if status in ("SUCCESS", "FAILURE"):
return data
time.sleep(3)

---------- 1. imagine 提交 ----------

prompt = "a cute steampunk cat, full body, high detail"
print("1️⃣ submitting imagine...")
r1 = _post("/submit/imagine", {"prompt": prompt})
print(r1)
task_id = r1["result"]
print(f" taskId = {task_id}")

---------- 2. 等待首次出图 ----------

print("2️⃣ waiting for imagine result...")
task_data = wait_task(task_id)
if task_data["status"] != "SUCCESS":
raise RuntimeError("imagine 失败", task_data)
image_url = task_data["imageUrl"]
print(f" imageUrl = {image_url}")

---------- 3. 先选图（以 U1 为例） ----------

buttons = task_data.get("buttons", [])
pprint(buttons)
u1_btn = next((b for b in buttons if b["label"] == "U1"), None)
if not u1_btn:
raise RuntimeError("未找到 U1 按钮，无法放大选图")
print("3️⃣ upscaling U1...")
r3 = _post("/submit/action", {"taskId": task_id, "customId": u1_btn["customId"]})
upscale_task_id = r3["result"]
upscale_data = wait_task(upscale_task_id) # 等待放大完成
if upscale_data["status"] != "SUCCESS":
raise RuntimeError("放大失败", upscale_data)

---------- 3. 找到局部重绘按钮 ----------

buttons = upscale_data.get("buttons", [])
pprint(buttons)
inpaint_btn = next((b for b in buttons if "Inpaint" in b["customId"]), None)
if not inpaint_btn:
raise RuntimeError("放大后仍无 Vary (Region)，图片分辨率可能仍不足")
custom_id = inpaint_btn["customId"]
print("4️⃣ triggering Vary (Region)...")
r4 = _post("/submit/action", {"taskId": upscale_task_id, "customId": custom_id})
modal_task_id = r4["result"]

---------- 4. 生成蒙版（简单画个白色矩形） ----------

print("4️⃣ creating mask...")
img_resp = session.get(image_url)
img = Image.open(BytesIO(img_resp.content)).convert("RGB")
W, H = img.size
""" 局部重绘，选择区域
mask = Image.new("L", (W, H), 0) # 全黑
draw = ImageDraw.Draw(mask)#
draw.rectangle([0, 0, W, H], fill=255)
"""
mask = Image.new("L", (W, H), 255) # 全局重绘
buf = BytesIO()
mask.save(buf, format="PNG")
mask_b64 = base64.b64encode(buf.getvalue()).decode()

---------- 5. 提交 modal（局部重绘） ----------

new_prompt = "a cute steampunk cat wearing futuristic VR goggles"
print("5️⃣ submitting modal...")
r5 = _post("/submit/modal", {
"taskId": modal_task_id,
"prompt": new_prompt,
"maskBase64": mask_b64
})
final_task_id = r5["result"]
print(f" final taskId = {final_task_id}")

---------- 6. 等待最终结果 ----------

print("6️⃣ waiting for inpaint result...")
final_data = wait_task(final_task_id)
if final_data["status"] != "SUCCESS":
raise RuntimeError("inpaint 失败", final_data)
print("✅ Done! 最终图片地址：", final_data["imageUrl"])

接口使用示例 (局部重绘)

---------- 工具函数 ----------#

---------- 1. imagine 提交 ----------#

---------- 2. 等待首次出图 ----------#

---------- 3. 先选图（以 U1 为例） ----------#

---------- 3. 找到局部重绘按钮 ----------#

---------- 4. 生成蒙版（简单画个白色矩形） ----------#

---------- 5. 提交 modal（局部重绘） ----------#

---------- 6. 等待最终结果 ----------#