Design Image Studio

__pycache__/
*.pyc
.DS_Store

MIT License

Copyright (c) 2026 kangarooking

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

volcengine-python-sdk[ark]

#!/usr/bin/env python3
"""
Design-aware wrapper for Volcengine image generation.

This script acts as a compiler:
1. User brief -> design reasoning
2. Design reasoning -> compiled brief
3. Compiled brief -> final image prompt
4. Final image prompt -> Volcengine generation
"""

from __future__ import annotations

import argparse
import json
import shlex
import subprocess
import sys
from pathlib import Path


TASK_KEYWORDS = {
    "poster": ["海报", "poster", "封面", "主视觉", "campaign", "kv"],
    "product": ["商品", "产品", "电商", "product", "hero shot", "广告图", "首图"],
    "ppt": ["ppt", "幻灯片", "演示", "deck", "slide", "配图", "章节页"],
    "infographic": ["信息图", "infographic", "流程图", "结构图", "总结图", "overview"],
    "teaching": ["教学", "演示图", "讲解", "培训", "课件", "步骤图", "demo"],
}

DEFAULT_ASPECT = {
    "poster": "3:4",
    "product": "1:1",
    "ppt": "16:9",
    "infographic": "4:3",
    "teaching": "16:9",
}

QUALITY_MODEL = {
    "draft": "doubao-seedream-4-0-250828",
    "final": "doubao-seedream-5-0-lite-260128",
    "premium": "doubao-seedream-4-5-251128",
}

ASPECT_SIZES = {
    "1:1": {"2K": "2048x2048", "3K": "3072x3072"},
    "3:4": {"2K": "1728x2304", "3K": "2592x3456"},
    "4:3": {"2K": "2304x1728", "3K": "3456x2592"},
    "16:9": {"2K": "2848x1600", "3K": "4096x2304"},
    "9:16": {"2K": "1600x2848", "3K": "2304x4096"},
    "3:2": {"2K": "2496x1664", "3K": "3744x2496"},
    "2:3": {"2K": "1664x2496", "3K": "2496x3744"},
    "4:5": {"2K": "1840x2304", "3K": "2760x3456"},
}

CORE_SOURCE_SECTIONS = [
    "## Your workflow",
    "## Output creation guidelines",
    "### How to do design work",
    "## Content Guidelines",
    "**Do not add filler content.**",
    "**Create a system up front:**",
    "**Avoid AI slop tropes:**",
]

TASK_SOURCE_SECTIONS = {
    "poster": [
        "### How to do design work",
        "## Content Guidelines",
        "**Do not add filler content.**",
    ],
    "product": [
        "## Output creation guidelines",
        "## Content Guidelines",
        "**Avoid AI slop tropes:**",
    ],
    "ppt": [
        "## Content Guidelines",
        "**Create a system up front:**",
        "**Use appropriate scales:**",
    ],
    "infographic": [
        "## Content Guidelines",
        "**Do not add filler content.**",
        "**Avoid AI slop tropes:**",
    ],
    "teaching": [
        "## Content Guidelines",
        "**Do not add filler content.**",
        "**Use appropriate scales:**",
    ],
}

DIRECTION_PROFILES = {
    "conservative": {
        "style_bias": "restrained, cleaner, corporate, polished, lower-risk",
        "energy_bias": "controlled, composed, authoritative",
        "composition_bias": "cleaner geometry, more negative space, lower background complexity",
        "palette_bias": "restricted palette with restrained accents",
        "detail_bias": "cleaner surfaces, less decorative detail",
    },
    "balanced": {
        "style_bias": "premium editorial, contemporary, polished, commercially strong",
        "energy_bias": "energetic but disciplined, confident, professional",
        "composition_bias": "clear hierarchy, dynamic but stable layout, deliberate focal contrast",
        "palette_bias": "premium neutrals with one strong accent family",
        "detail_bias": "hero-detail emphasis with a restrained background",
    },
    "bold": {
        "style_bias": "larger scale, more dramatic, more surprising, high-contrast",
        "energy_bias": "ambitious, high-energy, assertive, vivid",
        "composition_bias": "bolder crop, stronger scale contrast, more motion cues, still disciplined",
        "palette_bias": "higher-contrast palette with strong accent energy",
        "detail_bias": "high-impact hero detail, expressive texture, controlled spectacle",
    },
}

GLOBAL_DIRECTIVES = [
    "Start from purpose, audience, and channel rather than surface-level adjectives.",
    "Create a coherent visual system before detailing the image.",
    "Use one clear hero idea and preserve obvious hierarchy.",
    "Treat negative space and safe zones as design decisions, not leftover space.",
    "Avoid filler content, decorative noise, and meaningless visual data.",
    "Respect existing brand or context when available; if not, still commit to a clear direction.",
    "Avoid AI-slop tropes such as random HUD overlays, generic fog, empty gradients, and scattered floating debris.",
]

TASK_PROFILES = {
    "poster": {
        "communication_goal": "attract and persuade quickly in a campaign or recruitment context",
        "hero_strategy": "one dominant hero visual or symbolic concept, never a collage of equal-weight objects",
        "safe_zone": "reserve a clean, obvious text-safe zone in the upper third or one side for headline and CTA copy",
        "lighting": "crisp premium lighting with controlled contrast and energetic highlights",
        "palette": "restricted premium palette with disciplined neutrals and one energetic accent family",
        "detail_density": "hero-rich detail with restrained background complexity",
        "base_style": "editorial campaign key visual",
        "task_constraints": [
            "make the result feel campaign-ready rather than generic AI art",
            "prioritize focal hierarchy and typography-safe composition",
        ],
        "task_avoid": [
            "poster text rendered directly into the image unless explicitly requested",
            "random marketing icons and decorative interface fragments",
        ],
    },
    "product": {
        "communication_goal": "make the product feel desirable, premium, and commercially credible",
        "hero_strategy": "the product is the undisputed focal point with clear silhouette and edge readability",
        "safe_zone": "keep surrounding space supportive and uncluttered so the product remains dominant",
        "lighting": "commercial lighting that reveals material, finish, and shape without cheap reflections",
        "palette": "palette chosen to support product positioning rather than compete with the product",
        "detail_density": "high fidelity on the hero object, restrained props and background",
        "base_style": "high-end commercial product advertising",
        "task_constraints": [
            "preserve product proportions and perceived material integrity",
            "use props only when they reinforce product positioning",
        ],
        "task_avoid": [
            "oversized props stealing attention",
            "luxury claims paired with cheap lighting or noisy backgrounds",
        ],
    },
    "ppt": {
        "communication_goal": "support presentation storytelling with a clear, memorable visual metaphor",
        "hero_strategy": "a single strong metaphor or scene readable at presentation distance",
        "safe_zone": "reserve a large clean area for slide title and subtitle overlay",
        "lighting": "clean, legible lighting that supports shape readability over moodiness",
        "palette": "presentation-friendly palette with strong contrast and limited visual noise",
        "detail_density": "mid-detail image readable at a glance, not overloaded with tiny elements",
        "base_style": "presentation cover art",
        "task_constraints": [
            "favor readability from distance over excess detail",
            "keep enough room for future title placement",
        ],
        "task_avoid": [
            "ad-poster density",
            "tiny embedded text or fragile detail that disappears on slides",
        ],
    },
    "infographic": {
        "communication_goal": "communicate structure, grouping, and flow rather than literal dense data",
        "hero_strategy": "clear modular hierarchy with one dominant organizing principle",
        "safe_zone": "leave room for headings or labels without relying on the model to render tiny text",
        "lighting": "flat-to-controlled lighting that supports structure and clarity",
        "palette": "structured palette with clear grouping and low noise",
        "detail_density": "low-to-mid detail, with emphasis on grouping and directional logic",
        "base_style": "structured infographic-like visual system",
        "task_constraints": [
            "prioritize visual structure over fake data richness",
            "use symbolic clarity and modular composition",
        ],
        "task_avoid": [
            "tiny chart labels",
            "data slop, fake dashboards, and dense unreadable micro-details",
        ],
    },
    "teaching": {
        "communication_goal": "explain a process, comparison, or sequence with maximum clarity",
        "hero_strategy": "show the logic of the teaching point first, then add supporting visuals",
        "safe_zone": "keep panel or label areas simple and legible for later annotation",
        "lighting": "clear explanatory lighting, not overly cinematic",
        "palette": "clarity-first palette with simple grouping and controlled contrast",
        "detail_density": "mid-to-low detail with emphasis on readable sequence and large forms",
        "base_style": "instructional visual storytelling",
        "task_constraints": [
            "make sequence and cause-effect legible at first glance",
            "use big forms and obvious directional logic",
        ],
        "task_avoid": [
            "cinematic clutter that weakens explanation",
            "too many simultaneous steps in one frame",
        ],
    },
}


def detect_task(brief: str) -> str:
    text = brief.lower()
    scores = {}
    for task, keywords in TASK_KEYWORDS.items():
        scores[task] = sum(1 for keyword in keywords if keyword.lower() in text)
    best_task = max(scores, key=scores.get)
    return best_task if scores[best_task] > 0 else "poster"


def normalize_task(task: str, brief: str) -> str:
    return detect_task(brief) if task == "auto" else task


def choose_model(quality: str) -> str:
    return QUALITY_MODEL[quality]


def choose_size(task: str, aspect: str, quality: str) -> str:
    tier = "2K" if quality == "draft" or task in {"ppt", "infographic", "teaching"} else "3K"
    return ASPECT_SIZES.get(aspect, ASPECT_SIZES[DEFAULT_ASPECT[task]])[tier]


def join_phrases(items: list[str]) -> str:
    return "; ".join(item for item in items if item)


def unique_preserving_order(items: list[str]) -> list[str]:
    seen = set()
    ordered = []
    for item in items:
        if item not in seen:
            seen.add(item)
            ordered.append(item)
    return ordered


def make_design_reasoning(args: argparse.Namespace, task: str) -> dict:
    profile = TASK_PROFILES[task]
    direction = DIRECTION_PROFILES[args.direction]

    brand_strategy = (
        f"use the provided brand/context: {args.brand}"
        if args.brand
        else "no explicit brand provided; commit to one coherent visual system instead of averaging styles"
    )

    reference_strategy = (
        "use provided reference images to preserve consistency and context"
        if args.image
        else "no reference images provided; rely on the compiled visual system and brief"
    )

    source_sections = unique_preserving_order(CORE_SOURCE_SECTIONS + TASK_SOURCE_SECTIONS[task])

    visual_system = [
        f"base mode: {profile['base_style']}",
        f"direction bias: {direction['style_bias']}",
        f"energy bias: {direction['energy_bias']}",
        f"composition bias: {direction['composition_bias']}",
        f"palette bias: {direction['palette_bias']}",
    ]

    hierarchy_strategy = [
        "one clear hero idea",
        profile["hero_strategy"],
        "secondary elements must support the hero rather than compete with it",
        "background should create rhythm, not narrative confusion",
    ]

    anti_filler_rules = [
        "every element must earn its place",
        "do not add objects, labels, icons, or stats that do not strengthen the core message",
        "if the frame feels empty, solve with scale, crop, rhythm, or texture rather than random extra elements",
    ]

    anti_slop_rules = [
        "avoid generic AI clutter",
        "avoid random floating UI fragments or HUD overlays",
        "avoid generic gradient fog with no composition logic",
        "avoid cheap neon cyberpunk treatment unless explicitly requested",
        "avoid noisy micro-detail that weakens the hierarchy",
    ]

    if args.avoid:
        anti_slop_rules.append(args.avoid)

    return {
        "task": task,
        "direction": args.direction,
        "communication_goal": args.goal or profile["communication_goal"],
        "audience": args.audience or "broad professional audience",
        "channel": args.usage or task,
        "brief": args.brief.strip(),
        "brand_strategy": brand_strategy,
        "reference_strategy": reference_strategy,
        "visual_system": visual_system,
        "hierarchy_strategy": hierarchy_strategy,
        "safe_zone_strategy": args.safe_zone or profile["safe_zone"],
        "lighting_strategy": args.lighting or profile["lighting"],
        "palette_strategy": args.palette or profile["palette"],
        "detail_density": direction["detail_bias"] + "; " + profile["detail_density"],
        "style_direction": args.style or join_phrases(visual_system),
        "mood_direction": args.mood or direction["energy_bias"],
        "composition_logic": args.composition or direction["composition_bias"],
        "anti_filler_rules": anti_filler_rules,
        "anti_slop_rules": anti_slop_rules,
        "task_constraints": profile["task_constraints"] + ([args.constraints] if args.constraints else []),
        "task_avoid": profile["task_avoid"],
        "global_directives": GLOBAL_DIRECTIVES,
        "source_sections": source_sections,
        "primary_source_file": "references/claude-design-sys-prompt-full.txt",
    }


def compile_design_brief(reasoning: dict, aspect: str) -> dict:
    return {
        "task": reasoning["task"],
        "direction": reasoning["direction"],
        "brief": reasoning["brief"],
        "communication_goal": reasoning["communication_goal"],
        "audience": reasoning["audience"],
        "channel": reasoning["channel"],
        "brand_strategy": reasoning["brand_strategy"],
        "reference_strategy": reasoning["reference_strategy"],
        "visual_system": join_phrases(reasoning["visual_system"]),
        "hierarchy": join_phrases(reasoning["hierarchy_strategy"]),
        "composition": reasoning["composition_logic"],
        "safe_zone": reasoning["safe_zone_strategy"],
        "lighting": reasoning["lighting_strategy"],
        "palette": reasoning["palette_strategy"],
        "detail_density": reasoning["detail_density"],
        "style_direction": reasoning["style_direction"],
        "mood": reasoning["mood_direction"],
        "constraints": join_phrases(reasoning["task_constraints"]),
        "avoid": join_phrases(reasoning["anti_slop_rules"] + reasoning["task_avoid"]),
        "aspect": aspect,
        "source_sections": reasoning["source_sections"],
    }


def build_prompt(brief: dict) -> str:
    parts = [
        f"Create a {brief['task']} image for {brief['channel']} aimed at {brief['audience']}.",
        f"Treat this as a design-led visual solving this brief: {brief['brief']}.",
        f"Communication goal: {brief['communication_goal']}.",
        "Translate the brief into one strong hero concept rather than many equal-weight elements.",
        f"Brand and context strategy: {brief['brand_strategy']}.",
        f"Visual system: {brief['visual_system']}.",
        f"Hierarchy: {brief['hierarchy']}.",
        f"Composition: {brief['composition']}.",
        f"Safe zone: {brief['safe_zone']}.",
        f"Lighting: {brief['lighting']}.",
        f"Color strategy: {brief['palette']}.",
        f"Detail density: {brief['detail_density']}.",
        f"Style direction: {brief['style_direction']}.",
        f"Mood: {brief['mood']}.",
        f"Aspect ratio: {brief['aspect']}.",
        f"Important constraints: {brief['constraints']}.",
        f"Avoid: {brief['avoid']}.",
        "Emphasize strong hierarchy, intentional whitespace, disciplined background complexity, and polished professional finish.",
    ]
    return " ".join(parts)


def run_generation(
    args: argparse.Namespace,
    prompt: str,
    task: str,
    model: str,
    size: str,
) -> int:
    skill_root = Path(__file__).resolve().parent
    generate_script = skill_root / "generate.py"

    command = [
        sys.executable,
        str(generate_script),
        "--prompt",
        prompt,
        "--model",
        model,
        "--size",
        size,
        "--response-format",
        args.response_format,
        "--output-format",
        args.output_format,
    ]

    if args.output:
        command.extend(["--output", args.output])
    if args.output_dir:
        command.extend(["--output-dir", args.output_dir])
    if args.budget_limit is not None:
        command.extend(["--budget-limit", str(args.budget_limit)])
    if args.dry_run:
        command.append("--dry-run")
    if args.web_search:
        command.append("--web-search")
    if args.fast_mode:
        command.append("--fast-mode")
    if args.watermark:
        command.append("--watermark")
    if args.fallback_model:
        command.extend(["--fallback-model", args.fallback_model])
    if args.image:
        command.extend(["--image", *args.image])

    if task in {"teaching", "infographic"} and args.sequential:
        command.append("--sequential")
        command.extend(["--max-images", str(args.max_images)])

    print("\n[command]")
    print(" ".join(shlex.quote(part) for part in command))
    sys.stdout.flush()
    result = subprocess.run(command, check=False)
    return result.returncode


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Compile a design-system-driven brief into an image prompt and generate the image."
    )
    parser.add_argument("--task", default="auto", choices=["auto", "poster", "product", "ppt", "infographic", "teaching"])
    parser.add_argument("--brief", required=True, help="User brief or request")
    parser.add_argument("--audience", default=None, help="Target audience")
    parser.add_argument("--usage", default=None, help="Where the image will be used")
    parser.add_argument("--brand", default=None, help="Brand tone or reference context")
    parser.add_argument("--style", default=None, help="Preferred visual style override")
    parser.add_argument("--mood", default=None, help="Preferred mood override")
    parser.add_argument("--goal", default=None, help="Specific communication goal override")
    parser.add_argument("--composition", default=None, help="Composition override")
    parser.add_argument("--constraints", default=None, help="Must-have constraints")
    parser.add_argument("--avoid", default=None, help="Things to avoid")
    parser.add_argument("--aspect", default=None, help="Aspect ratio such as 1:1, 3:4, 16:9")
    parser.add_argument("--direction", default="balanced", choices=["conservative", "balanced", "bold"])
    parser.add_argument("--safe-zone", default=None, help="Safe-zone strategy override")
    parser.add_argument("--lighting", default=None, help="Lighting strategy override")
    parser.add_argument("--palette", default=None, help="Palette strategy override")
    parser.add_argument("--quality", default="final", choices=["draft", "final", "premium"])
    parser.add_argument("--model-override", default=None, help="Explicit model id override such as doubao-seedream-5-0-260128")
    parser.add_argument("--image", nargs="+", default=None, help="Reference image path(s) or URL(s)")
    parser.add_argument("--output", "-o", default=None, help="Output filename")
    parser.add_argument("--output-dir", default=None, help="Output directory")
    parser.add_argument("--output-format", default="png", choices=["png", "jpeg"])
    parser.add_argument("--response-format", default="b64_json", choices=["url", "b64_json"])
    parser.add_argument("--prompt-only", action="store_true", help="Only print design reasoning, compiled brief, and the final prompt")
    parser.add_argument("--budget-limit", type=float, default=None, help="Optional spend guard in CNY")
    parser.add_argument("--fallback-model", default=None, help="Fallback model id")
    parser.add_argument("--dry-run", action="store_true", help="Estimate cost without generating")
    parser.add_argument("--web-search", action="store_true", help="Enable web search if the selected model supports it")
    parser.add_argument("--fast-mode", action="store_true", help="Enable fast mode for compatible models")
    parser.add_argument("--watermark", action="store_true", help="Enable output watermark")
    parser.add_argument("--sequential", action="store_true", help="Generate sequential images for explanatory series")
    parser.add_argument("--max-images", type=int, default=4, help="Max images for sequential generation")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    task = normalize_task(args.task, args.brief)
    args.aspect = args.aspect or DEFAULT_ASPECT[task]
    model = args.model_override or choose_model(args.quality)
    size = choose_size(task, args.aspect, args.quality)

    design_reasoning = make_design_reasoning(args, task)
    compiled_brief = compile_design_brief(design_reasoning, args.aspect)
    prompt = build_prompt(compiled_brief)

    print("[design_reasoning]")
    print(json.dumps(design_reasoning, ensure_ascii=False, indent=2))
    print("\n[compiled_brief]")
    print(json.dumps(compiled_brief, ensure_ascii=False, indent=2))
    print("\n[prompt]")
    print(prompt)
    print(f"\n[settings]\nmodel={model}\nsize={size}\naspect={args.aspect}\ndirection={args.direction}")
    sys.stdout.flush()

    if args.prompt_only:
        return 0

    return run_generation(args, prompt, task, model, size)


if __name__ == "__main__":
    raise SystemExit(main())

#!/usr/bin/env python3
"""
Volcengine Seedream Image Generation Script
============================================
Features:
  - Text-to-image, image-to-image, multi-image fusion, sequential generation
  - Automatic retry with exponential backoff and jitter
  - Model fallback on repeated failures
  - Cost estimation, budget limiting, and usage logging
  - Streaming support for sequential images

Usage:
  python generate.py --prompt "A sunset over mountains" --model auto --size 2K
  python generate.py --prompt "Edit background" --image photo.png --size 2K
  python generate.py --prompt "Storyboard" --sequential --max-images 4
"""

import argparse
import base64
import json
import os
import random
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Optional

# ── Configuration ─────────────────────────────────────────────

MODELS = {
    "doubao-seedream-5-0-260128": {
        "name": "Seedream 5.0",
        "tier": "high",
        "resolutions": ["2K", "3K"],
        "formats": ["png", "jpeg"],
        "cost_per_image": {"1K": 0.04, "2K": 0.06, "3K": 0.10, "4K": 0.14},
    },
    "doubao-seedream-5-0-lite-260128": {
        "name": "Seedream 5.0 lite",
        "tier": "mid",
        "resolutions": ["2K", "3K"],
        "formats": ["png", "jpeg"],
        "cost_per_image": {"1K": 0.02, "2K": 0.04, "3K": 0.06, "4K": 0.08},
    },
    "doubao-seedream-4-5-251128": {
        "name": "Seedream 4.5",
        "tier": "high",
        "resolutions": ["2K", "4K"],
        "formats": ["jpeg"],
        "cost_per_image": {"1K": 0.04, "2K": 0.06, "3K": 0.08, "4K": 0.12},
    },
    "doubao-seedream-4-0-250828": {
        "name": "Seedream 4.0",
        "tier": "low",
        "resolutions": ["1K", "2K", "4K"],
        "formats": ["jpeg"],
        "cost_per_image": {"1K": 0.02, "2K": 0.03, "3K": 0.05, "4K": 0.08},
    },
}

# Resolution presets → approximate pixel counts
RESOLUTION_MAP = {
    "1K": 921600,
    "2K": 3686400,
    "3K": 9437184,
    "4K": 16777216,
}

COST_LOG_DEFAULT = os.path.expanduser("~/.volcengine-image-costs.json")

# ── Cost Management ───────────────────────────────────────────

def resolve_resolution(size_arg: str, model_id: str) -> str:
    """Resolve size argument to a resolution tier for cost estimation."""
    if size_arg in RESOLUTION_MAP:
        return size_arg
    # Parse pixel dimensions like "2048x2048"
    if "x" in size_arg.lower():
        parts = size_arg.lower().split("x")
        try:
            pixels = int(parts[0]) * int(parts[1])
            return min(
                RESOLUTION_MAP,
                key=lambda tier: abs(RESOLUTION_MAP[tier] - pixels),
            )
        except (ValueError, IndexError):
            pass
    return "2K"  # default


def estimate_cost(model_id: str, size: str, num_images: int = 1) -> float:
    """Estimate cost in CNY for a generation request."""
    model_info = MODELS.get(model_id, {})
    cost_table = model_info.get("cost_per_image", {})
    resolution = resolve_resolution(size, model_id)
    per_image = cost_table.get(resolution, 0.04)  # default estimate
    return per_image * num_images


def check_budget(cost: float, budget_limit: Optional[float], cost_log_path: str) -> bool:
    """Check if this generation would exceed the budget limit."""
    if budget_limit is None:
        return True

    cumulative = 0.0
    if os.path.exists(cost_log_path):
        try:
            with open(cost_log_path, "r") as f:
                log = json.load(f)
                cumulative = sum(e.get("cost", 0) for e in log.get("entries", []))
        except (json.JSONDecodeError, IOError):
            pass

    if cumulative + cost > budget_limit:
        print(f"⚠ Budget limit: cumulative ¥{cumulative:.4f} + estimated ¥{cost:.4f} "
              f"> limit ¥{budget_limit:.4f}", file=sys.stderr)
        return False
    return True


def log_cost(model_id: str, size: str, prompt_len: int, cost: float,
             success: bool, error: str, cost_log_path: str):
    """Append a cost entry to the JSON log."""
    log = {"entries": []}
    if os.path.exists(cost_log_path):
        try:
            with open(cost_log_path, "r") as f:
                log = json.load(f)
        except (json.JSONDecodeError, IOError):
            log = {"entries": []}

    log["entries"].append({
        "timestamp": datetime.now().isoformat(),
        "model": model_id,
        "resolution": size,
        "prompt_length": prompt_len,
        "estimated_cost_cny": round(cost, 4),
        "success": success,
        "error": error,
    })

    os.makedirs(os.path.dirname(cost_log_path) or ".", exist_ok=True)
    with open(cost_log_path, "w") as f:
        json.dump(log, f, indent=2, ensure_ascii=False)


# ── Retry Logic ───────────────────────────────────────────────

class RetryExhausted(Exception):
    """All retry attempts failed."""
    pass


class BudgetExceeded(Exception):
    """Generation would exceed budget."""
    pass


def classify_error(status_code: int, error_msg: str) -> str:
    """Classify an API error to determine retry strategy.

    Returns one of:
      - "retry"        → transient, safe to retry
      - "auth"         → authentication error, must stop
      - "bad_request"  → invalid params, must fix
      - "content_policy" → safety filter, must revise prompt
      - "unknown"      → unclassified
    """
    if status_code in (401, 403):
        return "auth"
    if status_code == 400:
        if "content" in error_msg.lower() or "safety" in error_msg.lower():
            return "content_policy"
        return "bad_request"
    if status_code == 429:
        return "retry"
    if status_code >= 500:
        return "retry"
    if "timeout" in error_msg.lower() or "connection" in error_msg.lower():
        return "retry"
    return "unknown"


def retry_with_backoff(func, max_retries=3, initial_delay=1.0,
                        multiplier=2.0, on_fallback=None):
    """Execute func with exponential backoff + jitter on retriable errors.

    Args:
        func: Callable that performs the API call. Should raise on failure.
        max_retries: Maximum retry attempts.
        initial_delay: First retry delay in seconds.
        multiplier: Backoff multiplier between retries.
        on_fallback: Optional callable(model_id) to switch to a fallback model.

    Returns:
        The result of func().

    Raises:
        RetryExhausted: All retries failed.
    """
    last_error = None

    for attempt in range(max_retries + 1):
        try:
            return func()
        except Exception as e:
            last_error = e

            # Extract status code from various SDK exception formats
            status_code = getattr(e, "status_code", None) or 0
            error_msg = str(e)

            error_type = classify_error(status_code, error_msg)

            # Non-retriable errors: fail immediately
            if error_type == "auth":
                print(f"✘ Authentication error: {error_msg}", file=sys.stderr)
                raise
            if error_type == "bad_request":
                print(f"✘ Bad request: {error_msg}", file=sys.stderr)
                raise
            if error_type == "content_policy":
                print(f"✘ Content policy violation: {error_msg}", file=sys.stderr)
                print("  Suggestion: Revise your prompt to avoid sensitive content.",
                      file=sys.stderr)
                raise

            # Retriable: attempt retry or fallback
            if attempt < max_retries:
                delay = initial_delay * (multiplier ** attempt)
                jitter = random.uniform(0, delay * 0.3)
                total_delay = delay + jitter

                print(f"⟳ Attempt {attempt + 1}/{max_retries} failed "
                      f"({error_type}): {error_msg[:100]}", file=sys.stderr)
                print(f"  Retrying in {total_delay:.1f}s...", file=sys.stderr)

                # On last retry attempt, try fallback model if provided
                if attempt == max_retries - 1 and on_fallback:
                    fallback_msg = on_fallback()
                    if fallback_msg:
                        print(f"  ↓ Falling back to: {fallback_msg}", file=sys.stderr)

                time.sleep(total_delay)
            else:
                print(f"✘ All {max_retries} retries exhausted.", file=sys.stderr)

    raise RetryExhausted(f"Failed after {max_retries} retries. Last error: {last_error}")


# ── Image Helpers ─────────────────────────────────────────────

def image_to_data_uri(image_path: str) -> str:
    """Convert a local image file to a data URI for the API."""
    ext = Path(image_path).suffix.lower().lstrip(".")
    if ext == "jpg":
        ext = "jpeg"
    with open(image_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("utf-8")
    return f"data:image/{ext};base64,{b64}"


def download_image(url: str, output_path: str):
    """Download an image from URL to local file."""
    import urllib.request
    urllib.request.urlretrieve(url, output_path)


# ── Main Generation Logic ────────────────────────────────────

def select_model(model_arg: str) -> str:
    """Resolve 'auto' or shorthand to a full model ID."""
    if model_arg == "auto":
        return "doubao-seedream-5-0-lite-260128"  # best value default

    # Allow short names
    aliases = {
        "5.0": "doubao-seedream-5-0-260128",
        "5.0-lite": "doubao-seedream-5-0-lite-260128",
        "4.5": "doubao-seedream-4-5-251128",
        "4.0": "doubao-seedream-4-0-250828",
    }
    return aliases.get(model_arg, model_arg)


def run_generation(args):
    """Execute the image generation with all features."""
    # ── Resolve model ──
    model_id = select_model(args.model)
    fallback_model = select_model(args.fallback_model) if args.fallback_model else None

    if model_id not in MODELS:
        print(f"✘ Unknown model: {model_id}", file=sys.stderr)
        print(f"  Available: {', '.join(MODELS.keys())}", file=sys.stderr)
        sys.exit(1)

    # ── Resolve images ──
    images = None
    if args.image:
        images = []
        for img_path in args.image:
            if img_path.startswith(("http://", "https://")):
                images.append(img_path)
            elif os.path.isfile(img_path):
                images.append(image_to_data_uri(img_path))
            else:
                print(f"✘ Image not found: {img_path}", file=sys.stderr)
                sys.exit(1)

    # ── Estimate cost ──
    num_images = args.max_images if args.sequential else 1
    estimated_cost = estimate_cost(model_id, args.size, num_images)

    print(f"Model: {MODELS[model_id]['name']} ({model_id})")
    print(f"Resolution: {args.size}")
    print(f"Images: {num_images}")
    print(f"Estimated cost: ¥{estimated_cost:.4f}")

    if args.dry_run:
        print("\n[Dry run] No API call made.")
        return

    # ── Budget check ──
    if not check_budget(estimated_cost, args.budget_limit, args.cost_log):
        print("✘ Aborted: would exceed budget limit.", file=sys.stderr)
        sys.exit(1)

    # ── Prepare API client ──
    try:
        from volcenginesdkarkruntime import Ark
    except ImportError:
        print("✘ volcengine-python-sdk not installed.", file=sys.stderr)
        print("  Run: pip install 'volcengine-python-sdk[ark]'", file=sys.stderr)
        sys.exit(1)

    api_key = os.environ.get("ARK_API_KEY")
    if not api_key:
        print("✘ ARK_API_KEY environment variable not set.", file=sys.stderr)
        print("  Get your key at: https://console.volcengine.com/ark/region:ark+cn-beijing/apikey",
              file=sys.stderr)
        sys.exit(1)

    client = Ark(
        base_url="https://ark.cn-beijing.volces.com/api/v3",
        api_key=api_key,
    )

    # ── Build request kwargs ──
    gen_kwargs = {
        "model": model_id,
        "prompt": args.prompt,
        "size": args.size,
        "response_format": args.response_format,
        "watermark": args.watermark,
    }

    # Seedream 4.5 rejects the `output_format` parameter entirely.
    if model_id != "doubao-seedream-4-5-251128":
        gen_kwargs["output_format"] = args.output_format

    if images:
        gen_kwargs["image"] = images if len(images) > 1 else images[0]

    if args.sequential:
        from volcenginesdkarkruntime.types.images.images import (
            SequentialImageGenerationOptions,
        )
        gen_kwargs["sequential_image_generation"] = "auto"
        gen_kwargs["sequential_image_generation_options"] = (
            SequentialImageGenerationOptions(max_images=args.max_images)
        )

    if args.web_search:
        gen_kwargs["tools"] = [{"type": "web_search"}]

    if args.fast_mode:
        gen_kwargs["optimize_prompt_options"] = {"mode": "fast"}

    # ── Execute with retry ──
    current_model = model_id

    def do_call():
        gen_kwargs["model"] = current_model
        if current_model == "doubao-seedream-4-5-251128":
            gen_kwargs.pop("output_format", None)
        else:
            gen_kwargs["output_format"] = args.output_format
        return client.images.generate(**gen_kwargs)

    def try_fallback():
        nonlocal current_model
        if fallback_model and fallback_model != current_model:
            current_model = fallback_model
            return MODELS[fallback_model]["name"]
        return None

    try:
        result = retry_with_backoff(
            do_call,
            max_retries=args.max_retries,
            initial_delay=args.retry_delay,
            multiplier=args.retry_multiplier,
            on_fallback=try_fallback if fallback_model else None,
        )
    except RetryExhausted as e:
        log_cost(model_id, args.size, len(args.prompt), estimated_cost,
                 False, str(e), args.cost_log)
        print(f"\n✘ Generation failed: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        log_cost(model_id, args.size, len(args.prompt), estimated_cost,
                 False, str(e), args.cost_log)
        print(f"\n✘ Generation failed: {e}", file=sys.stderr)
        sys.exit(1)

    # ── Process results ──
    output_dir = args.output_dir or "."
    os.makedirs(output_dir, exist_ok=True)

    saved_files = []
    for i, img_data in enumerate(result.data):
        if args.response_format == "url" and hasattr(img_data, "url") and img_data.url:
            if args.output:
                filename = args.output if num_images == 1 else f"{Path(args.output).stem}_{i+1}{Path(args.output).suffix}"
            else:
                ext = "png" if args.output_format == "png" else "jpg"
                filename = f"generated_{int(time.time())}_{i+1}.{ext}"
            filepath = os.path.join(output_dir, filename)
            download_image(img_data.url, filepath)
            saved_files.append(filepath)
            print(f"  ✓ Saved: {filepath}")
        elif args.response_format == "b64_json" and hasattr(img_data, "b64_json"):
            if args.output:
                filename = args.output if num_images == 1 else f"{Path(args.output).stem}_{i+1}{Path(args.output).suffix}"
            else:
                ext = "png" if args.output_format == "png" else "jpg"
                filename = f"generated_{int(time.time())}_{i+1}.{ext}"
            filepath = os.path.join(output_dir, filename)
            with open(filepath, "wb") as f:
                f.write(base64.b64decode(img_data.b64_json))
            saved_files.append(filepath)
            print(f"  ✓ Saved: {filepath}")

    # ── Log cost ──
    actual_cost = estimated_cost  # Use estimate (API doesn't return exact cost)
    log_cost(current_model, args.size, len(args.prompt), actual_cost,
             True, "", args.cost_log)

    total = sum(e.get("estimated_cost_cny", 0) for e in
                json.load(open(args.cost_log)).get("entries", [])) \
            if os.path.exists(args.cost_log) else actual_cost

    print(f"\n  Cost this request: ~¥{actual_cost:.4f}")
    print(f"  Files saved: {len(saved_files)}")

    return saved_files


# ── CLI ───────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Generate images via Volcengine Seedream API",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    # ── Required ──
    parser.add_argument("--prompt", required=True, help="Text prompt for image generation")

    # ── Model & Output ──
    parser.add_argument("--model", default="auto",
                        help="Model: auto, 5.0, 5.0-lite, 4.5, 4.0, or full model ID")
    parser.add_argument("--fallback-model", default=None,
                        help="Fallback model if primary fails after retries")
    parser.add_argument("--size", default="2K",
                        help="Resolution: 1K/2K/3K/4K or WxH (e.g. 2048x2048)")
    parser.add_argument("--output", "-o", default=None, help="Output filename")
    parser.add_argument("--output-dir", default=None, help="Output directory")
    parser.add_argument("--output-format", default="png",
                        choices=["png", "jpeg"], help="Output image format")
    parser.add_argument("--response-format", default="url",
                        choices=["url", "b64_json"], help="API response format")

    # ── Input Images ──
    parser.add_argument("--image", nargs="+", default=None,
                        help="Input image path(s) or URL(s) for image-to-image")

    # ── Sequential Generation ──
    parser.add_argument("--sequential", action="store_true",
                        help="Enable sequential image generation")
    parser.add_argument("--max-images", type=int, default=4,
                        help="Max images for sequential generation (default: 4)")

    # ── Advanced Features ──
    parser.add_argument("--web-search", action="store_true",
                        help="Enable web search (5.0 lite only)")
    parser.add_argument("--fast-mode", action="store_true",
                        help="Use fast prompt optimization (4.0 only)")
    parser.add_argument("--watermark", action="store_true",
                        help="Add AI-generated watermark")

    # ── Retry & Error Handling ──
    parser.add_argument("--max-retries", type=int, default=3,
                        help="Max retry attempts (default: 3)")
    parser.add_argument("--retry-delay", type=float, default=1.0,
                        help="Initial retry delay in seconds (default: 1)")
    parser.add_argument("--retry-multiplier", type=float, default=2.0,
                        help="Backoff multiplier (default: 2)")

    # ── Cost Management ──
    parser.add_argument("--budget-limit", type=float, default=None,
                        help="Max cumulative spend in CNY")
    parser.add_argument("--cost-log", default=COST_LOG_DEFAULT,
                        help=f"Cost log JSON path (default: {COST_LOG_DEFAULT})")
    parser.add_argument("--dry-run", action="store_true",
                        help="Estimate cost without calling API")

    args = parser.parse_args()

    run_generation(args)


if __name__ == "__main__":
    main()