Files
workbuddy_skills/draw-ui/scripts/generate_image.py
T
admin 237a09a5b0 修:draw-ui / oil-motion 原被当子模块指针收录 ⇒ 改为正常文件入库(两份内容原先对别人是空的)
一、问题(本轮实测)
`draw-ui` 与 `oil-motion` 目录里**各自带一个内嵌 `.git`** ⇒ 上一次提交把它们记成了 **gitlink(子模块指针)**
⇒ 仓库里只存了一个不属于任何远端的 commit id,**别人克隆下来这两份是空的** ✗(`git status` 显示 ` m draw-ui` / ` m oil-motion` = 子模块内容有改动)。

二、处置(可回退)
· 把两处的 `.git` **挪走**(⛔ 不是删除)⇒ `归档/内嵌git-20261008/{draw-ui,oil-motion}.git`;
· `git rm --cached` 掉那两个 gitlink,再 `git add` 两个目录 ⇒ **按正常文件入库**(内容才真的进仓库)。

三、副作用(如实记)
挪走 `.git` 后,这两个技能**不能再原地 `git pull` 取上游更新**(要更新得重新拉一份覆盖);
如需恢复其本地仓库,把 `归档/内嵌git-20261008/` 里的 `.git` 挪回原处即可。
2026-10-08 22:29:52 +08:00

547 lines
21 KiB
Python

#!/usr/bin/env python3
from __future__ import annotations
import argparse
import base64
import binascii
import json
import mimetypes
import os
import re
import sys
import tempfile
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Iterable
DEFAULT_MODEL = os.getenv("DRAW_MODEL", "openai/gpt-image-2")
DEFAULT_PROVIDER = os.getenv("DRAW_PROVIDER", "zenmux")
DEFAULT_BASE_URL = os.getenv("ZENMUX_VERTEX_BASE_URL", "https://zenmux.ai/api/vertex-ai")
DEFAULT_OUTPUT_ROOT = Path.home() / ".local" / "share" / "draw" / "outputs"
DEFAULT_MIME = "image/png"
DEFAULT_CODEX_MODEL = os.getenv("DRAW_CODEX_MODEL", "gpt-5.6")
def _read_env_value(path: Path, key: str) -> str:
try:
text = path.read_text(encoding="utf-8")
except FileNotFoundError:
return ""
for line in text.splitlines():
stripped = line.strip()
if not stripped or stripped.startswith("#") or "=" not in stripped:
continue
name, value = stripped.split("=", 1)
if name.strip() != key:
continue
value = value.strip().strip('"').strip("'")
return value
return ""
def resolve_api_key() -> str:
env = os.getenv("ZENMUX_API_KEY", "").strip()
if env:
return env
cwd = Path.cwd().resolve()
for directory in [cwd, *cwd.parents]:
found = _read_env_value(directory / ".env.local", "ZENMUX_API_KEY")
if found:
return found
config_path = Path.home() / ".config" / "see" / "api_key"
try:
return config_path.read_text(encoding="utf-8").strip()
except FileNotFoundError:
return ""
def sanitize_name(value: str, fallback: str = "image") -> str:
value = value.strip()
value = re.sub(r"[\\/:*?\"<>|]+", "-", value)
value = re.sub(r"\s+", "-", value)
value = re.sub(r"-+", "-", value).strip("-_.")
if not value:
return fallback
return value[:80]
def build_output_path(*, output_arg: str, image_type: str, topic: str, explicit_name: str, ext: str) -> Path:
if output_arg:
out = Path(output_arg).expanduser().resolve()
if out.suffix:
return out
return out.with_suffix(ext)
now = datetime.now()
day_dir = DEFAULT_OUTPUT_ROOT / now.strftime("%Y-%m-%d")
day_dir.mkdir(parents=True, exist_ok=True)
base_name = sanitize_name(explicit_name or topic, fallback=image_type)
return day_dir / f"{now.strftime('%Y%m%d-%H%M%S')}__{image_type}__{base_name}{ext}"
def metadata_path_for(image_path: Path) -> Path:
return image_path.with_suffix(image_path.suffix + ".json")
def _output_paths_for_check(output_path: Path) -> list[Path]:
"""Return every path that a writer may choose before its MIME type is known."""
output_path = output_path.resolve()
if output_path.suffix:
return [output_path]
# The current writers default to PNG, while render_response may infer another
# image extension from the model response. Include existing siblings so an
# unknown future MIME type cannot silently replace one of them.
paths = [output_path.with_suffix(".png")]
pattern = f"{output_path.name}.*"
paths.extend(path for path in output_path.parent.glob(pattern) if path.is_file())
return list(dict.fromkeys(paths))
def ensure_output_available(output_path: Path) -> None:
"""Reject an output or its metadata before any remote generation starts."""
conflicts: list[Path] = []
for candidate in _output_paths_for_check(output_path):
if candidate.exists():
conflicts.append(candidate)
metadata_path = metadata_path_for(candidate)
if metadata_path.exists():
conflicts.append(metadata_path)
if conflicts:
paths = ", ".join(str(path) for path in dict.fromkeys(conflicts))
raise FileExistsError(f"输出或 metadata 已存在,拒绝覆盖:{paths}")
def _write_new_bytes(path: Path, data: bytes) -> None:
"""Create a file without ever replacing an existing output."""
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("xb") as handle:
handle.write(data)
def guess_extension(mime_type: str | None) -> str:
if not mime_type:
return ".png"
guessed = mimetypes.guess_extension(mime_type)
if guessed == ".jpe":
return ".jpg"
return guessed or ".png"
# Type only controls aspect ratio, prompt is fully controlled by caller
ASPECT_RATIOS = {
"ultrawide": "21:9",
"wide": "16:9",
"square": "1:1",
"portrait": "3:4",
"classic": "4:3",
}
CODEX_SIZE_PRESETS = {
"ultrawide": "1536x640",
"wide": "1536x864",
"classic": "1024x768",
"square": "1024x1024",
"portrait": "768x1024",
}
MODE_PROMPTS = {
"normal": "",
"replicate": (
"Use the reference image as the primary visual source. Recreate the UI screen as closely as possible. "
"Preserve layout, spacing, typography hierarchy, colors, shadows, border radius, icon style, density, "
"and the relative position of every major element. Do not redesign unless the prompt explicitly asks for changes. "
"If text is unreadable, preserve its visual length, alignment, and hierarchy. Output only the clean UI mockup, "
"with no browser chrome, watermark, annotations, or surrounding device frame."
),
"frame-lock": (
"Use the first reference image as a locked application frame. Preserve the sidebar, top navigation, brand area, "
"and persistent chrome as closely as possible. Redesign or generate only the content area requested by the prompt. "
"Keep the result as a clean full-screen UI mockup with no browser chrome or watermark."
),
"asset-redraw": (
"Use the reference image to recreate only the requested visual asset as a clean standalone asset. Remove surrounding "
"UI, labels, browser chrome, mockup frames, and unrelated elements unless explicitly requested. Preserve the source "
"asset's proportions, material, color, and brand feel with high clarity and generous padding."
),
}
def effective_prompt(prompt: str, mode: str) -> str:
mode_prompt = MODE_PROMPTS.get(mode, "")
if not mode_prompt:
return prompt
return f"{mode_prompt}\n\nUser request:\n{prompt}"
def download_file(url: str, dest: Path, timeout: int = 120) -> None:
req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
with urllib.request.urlopen(req, timeout=timeout) as response:
dest.write_bytes(response.read())
def resolve_ref(raw: str, tmp_dir: Path) -> Path:
parsed = urllib.parse.urlparse(raw)
if parsed.scheme in {"http", "https"}:
suffix = Path(parsed.path).suffix or ".png"
dest = tmp_dir / f"ref-{len(list(tmp_dir.iterdir())) + 1}{suffix}"
download_file(raw, dest)
return dest
path = Path(raw).expanduser().resolve()
if not path.exists():
raise FileNotFoundError(f"Reference not found: {raw}")
return path
def load_genai():
try:
from google import genai
from google.genai import types
except ModuleNotFoundError as exc:
raise SystemExit(
"[ERROR] Missing dependency `google-genai`. Re-run via scripts/ask_draw.sh so it can auto-install it."
) from exc
return genai, types
def build_contents(*, prompt: str, types, refs: Iterable[Path]):
parts = [types.Part.from_text(text=prompt)]
for ref in refs:
mime = mimetypes.guess_type(ref.name)[0] or DEFAULT_MIME
parts.append(types.Part.from_bytes(data=ref.read_bytes(), mime_type=mime))
return parts
def extract_parts(response) -> list:
if getattr(response, "parts", None):
return list(response.parts)
candidates = getattr(response, "candidates", None) or []
parts = []
for candidate in candidates:
content = getattr(candidate, "content", None)
if content and getattr(content, "parts", None):
parts.extend(content.parts)
return parts
def render_response(*, response, output_path: Path) -> tuple[str, str]:
text_parts: list[str] = []
image_written = False
image_mime = DEFAULT_MIME
pending_bytes: bytes | None = None
for part in extract_parts(response):
text = getattr(part, "text", None)
if text:
text_parts.append(text.strip())
inline_data = getattr(part, "inline_data", None)
if inline_data:
data = inline_data.data
if isinstance(data, str):
pending_bytes = base64.b64decode(data)
else:
pending_bytes = data
image_mime = getattr(inline_data, "mime_type", None) or DEFAULT_MIME
if pending_bytes:
final_path = output_path
if not output_path.suffix:
final_path = output_path.with_suffix(guess_extension(image_mime))
_write_new_bytes(final_path, pending_bytes)
image_written = True
else:
final_path = output_path
if not image_written:
raise RuntimeError("Model returned no image data.")
return final_path.as_posix(), "\n".join([t for t in text_parts if t]).strip()
def resolve_codex_api_key() -> str:
key = os.getenv("OPENAI_IMAGE_API_KEY") or os.getenv("OPENAI_API_KEY")
return key.strip() if key else ""
def resolve_codex_base_url() -> str:
base_url = (os.getenv("OPENAI_IMAGE_BASE_URL") or "https://api.openai.com/v1").rstrip("/")
parsed = urllib.parse.urlparse(base_url)
if parsed.scheme and parsed.netloc and parsed.path in ("", "/"):
return f"{base_url}/v1"
return base_url
def resolve_codex_model(override: str = "") -> str:
return override.strip() or DEFAULT_CODEX_MODEL
def join_endpoint(base_url: str, endpoint: str) -> str:
base = base_url.rstrip("/")
endpoint = endpoint.lstrip("/")
if base.endswith("/v1") and endpoint.startswith("v1/"):
endpoint = endpoint[3:]
return f"{base}/{endpoint}"
def ref_to_input_image(path: Path) -> dict:
mime = mimetypes.guess_type(path.name)[0] or DEFAULT_MIME
encoded = base64.b64encode(path.read_bytes()).decode("ascii")
return {"type": "input_image", "image_url": f"data:{mime};base64,{encoded}"}
def looks_like_base64_image(value: str) -> bool:
if len(value) < 200:
return False
compact = value.strip()
if compact.startswith("data:image/"):
compact = compact.split(",", 1)[-1]
try:
head = base64.b64decode(compact[:256] + "==", validate=False)
except Exception:
return False
return head.startswith(b"\x89PNG") or head.startswith(b"\xff\xd8\xff") or head.startswith(b"RIFF")
def find_image_result_recursive(value: object) -> str | None:
if isinstance(value, dict):
value_type = value.get("type")
for key in ("result", "b64_json", "image_base64"):
item = value.get(key)
if isinstance(item, str) and (value_type == "image_generation_call" or looks_like_base64_image(item)):
return item.split(",", 1)[-1] if item.startswith("data:image/") else item
for item in value.values():
found = find_image_result_recursive(item)
if found:
return found
elif isinstance(value, list):
for item in value:
found = find_image_result_recursive(item)
if found:
return found
elif isinstance(value, str) and looks_like_base64_image(value):
return value.split(",", 1)[-1] if value.startswith("data:image/") else value
return None
def request_codex_image(
*, prompt: str, refs: list[Path], image_type: str, model: str, output_path: Path
) -> Path:
api_key = resolve_codex_api_key()
if not api_key:
raise RuntimeError("No OPENAI_IMAGE_API_KEY or OPENAI_API_KEY found.")
base_url = resolve_codex_base_url()
endpoint = join_endpoint(base_url, "responses")
content: list[dict] = [{"type": "input_text", "text": prompt}]
content.extend(ref_to_input_image(ref) for ref in refs)
payload = {
"model": model,
"instructions": "Use the image_generation tool to create exactly the requested image. Do not add extra text.",
"stream": False,
"store": False,
"input": [{"role": "user", "content": content}],
"tools": [{"type": "image_generation", "size": CODEX_SIZE_PRESETS.get(image_type, "1024x1024")}],
"tool_choice": "required",
}
request = urllib.request.Request(
endpoint,
data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"Accept": "application/json",
},
method="POST",
)
try:
with urllib.request.urlopen(request, timeout=600) as response:
raw = response.read().decode("utf-8")
except urllib.error.HTTPError as exc:
details = exc.read().decode("utf-8", errors="replace")
raise RuntimeError(f"HTTP {exc.code} from Codex image API:\n{details}") from exc
except urllib.error.URLError as exc:
raise RuntimeError(f"Could not reach Codex image API: {exc.reason}") from exc
try:
response_payload = json.loads(raw)
except json.JSONDecodeError as exc:
raise RuntimeError(f"Codex image API returned non-JSON response:\n{raw[:1500]}") from exc
image_b64 = find_image_result_recursive(response_payload)
if not image_b64:
raise RuntimeError("No image result found in Codex Responses API output.")
final_path = output_path if output_path.suffix else output_path.with_suffix(".png")
final_path.parent.mkdir(parents=True, exist_ok=True)
try:
image_bytes = base64.b64decode(image_b64, validate=True)
except (binascii.Error, ValueError, TypeError) as exc:
raise RuntimeError("Codex image API returned invalid base64 image data.") from exc
_write_new_bytes(final_path, image_bytes)
return final_path
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate UI images via ZenMux or a Codex OpenAI-compatible provider.")
parser.add_argument("--type", choices=sorted(ASPECT_RATIOS.keys()), default="wide", help="Aspect ratio preset.")
parser.add_argument("--prompt", required=True, help="Full prompt for image generation.")
parser.add_argument("--ref", action="append", default=[], help="Reference image path or URL (repeatable).")
parser.add_argument("--name", default="", help="Optional short output name.")
parser.add_argument("-o", "--output", default="", help="Output image path.")
parser.add_argument("--provider", choices=["zenmux", "codex"], default=DEFAULT_PROVIDER, help="Image backend.")
parser.add_argument("--mode", choices=sorted(MODE_PROMPTS.keys()), default="normal", help="UI prompt wrapper.")
parser.add_argument("--model", default="", help="Model override (defaults per provider).")
parser.add_argument("--base-url", default=DEFAULT_BASE_URL, help=argparse.SUPPRESS)
return parser.parse_args()
def _uses_generate_images_api(model: str) -> bool:
"""Models that require the generate_images / edit_image API instead of generate_content."""
return model.startswith("openai/")
def _run_generate_images(*, client, model: str, prompt: str, refs: list[Path], types) -> tuple[bytes, str]:
"""Call generate_images (or edit_image when refs are provided) and return (image_bytes, response_text)."""
if refs:
# Use edit_image with reference images
# First ref becomes the base image
base_image_path = refs[0]
base_mime = mimetypes.guess_type(base_image_path.name)[0] or DEFAULT_MIME
base_image = types.Image(image_bytes=base_image_path.read_bytes(), mime_type=base_mime)
reference_images = [
types.RawReferenceImage(reference_id=1, reference_image=base_image)
]
# Additional refs as extra references
for i, ref_path in enumerate(refs[1:], start=2):
ref_mime = mimetypes.guess_type(ref_path.name)[0] or DEFAULT_MIME
ref_img = types.Image(image_bytes=ref_path.read_bytes(), mime_type=ref_mime)
reference_images.append(
types.RawReferenceImage(reference_id=i, reference_image=ref_img)
)
response = client.models.edit_image(
model=model,
prompt=prompt,
reference_images=reference_images,
)
else:
response = client.models.generate_images(
model=model,
prompt=prompt,
)
generated = getattr(response, "generated_images", None)
if not generated:
raise RuntimeError("Model returned no generated images.")
image_obj = generated[0].image
image_bytes = getattr(image_obj, "image_bytes", None)
if image_bytes is None:
# Some versions expose .data as base64
raw = getattr(image_obj, "data", None)
if isinstance(raw, str):
image_bytes = base64.b64decode(raw)
elif isinstance(raw, bytes):
image_bytes = raw
if not image_bytes:
raise RuntimeError("Could not extract image bytes from generate_images response.")
return image_bytes, ""
def main() -> int:
args = parse_args()
aspect_ratio = ASPECT_RATIOS[args.type]
prompt = effective_prompt(args.prompt, args.mode)
model = resolve_codex_model(args.model) if args.provider == "codex" else (args.model or DEFAULT_MODEL)
output_path = build_output_path(
output_arg=args.output,
image_type=args.type,
topic=args.name or "image",
explicit_name=args.name,
ext=".png",
)
ensure_output_available(output_path)
with tempfile.TemporaryDirectory(prefix="draw-refs-") as tmp:
tmp_dir = Path(tmp)
refs = [resolve_ref(raw, tmp_dir) for raw in args.ref]
if args.provider == "codex":
final_path = request_codex_image(
prompt=prompt,
refs=refs,
image_type=args.type,
model=model,
output_path=output_path,
)
response_text = ""
else:
api_key = resolve_api_key()
if not api_key:
print(
"[ERROR] No ZENMUX_API_KEY found. Set it as env var, in .env.local, or in ~/.config/see/api_key",
file=sys.stderr,
)
return 1
genai, types = load_genai()
# OpenAI image models via ZenMux can take longer; bump timeout to 5 minutes.
timeout = 300 if _uses_generate_images_api(model) else 120
client = genai.Client(
api_key=api_key,
vertexai=True,
http_options=types.HttpOptions(api_version="v1", base_url=args.base_url, timeout=timeout * 1000),
)
if _uses_generate_images_api(model):
image_bytes, response_text = _run_generate_images(
client=client, model=model, prompt=prompt, refs=refs, types=types,
)
final_path = output_path if output_path.suffix else output_path.with_suffix(".png")
final_path.parent.mkdir(parents=True, exist_ok=True)
_write_new_bytes(final_path, image_bytes)
else:
response = client.models.generate_content(
model=model,
contents=build_contents(prompt=prompt, types=types, refs=refs),
config=types.GenerateContentConfig(
response_modalities=["TEXT", "IMAGE"],
image_config=types.ImageConfig(aspect_ratio=aspect_ratio),
),
)
final_path_str, response_text = render_response(response=response, output_path=output_path)
final_path = Path(final_path_str)
meta_path = metadata_path_for(final_path)
metadata = {
"created_at": datetime.now().isoformat(timespec="seconds"),
"type": args.type,
"aspect_ratio": aspect_ratio,
"prompt": prompt,
"raw_prompt": args.prompt,
"refs": [str(path) for path in refs],
"provider": args.provider,
"mode": args.mode,
"model": model,
"base_url": args.base_url if args.provider == "zenmux" else resolve_codex_base_url(),
"output_path": str(final_path),
"response_text": response_text,
}
meta_path.parent.mkdir(parents=True, exist_ok=True)
with meta_path.open("x", encoding="utf-8") as handle:
handle.write(json.dumps(metadata, ensure_ascii=False, indent=2))
print(f"output_path={final_path}")
print(f"metadata_path={meta_path}")
return 0
if __name__ == "__main__":
raise SystemExit(main())