返回 ppt-master
backend_fal.py
根目录 / skills / ppt-master / scripts / image_backends / backend_fal.py
1 #!/usr/bin/env python3
2 """
3 fal.ai image generation backend.
4
5 Configuration keys:
6 FAL_KEY / FAL_API_KEY (required)
7 FAL_BASE_URL (optional)
8 FAL_MODEL (optional; nano-banana-2 only)
9 """
10
11 import sys
12 from pathlib import Path
13
14 _SCRIPTS_DIR = Path(__file__).resolve().parents[1]
15 if str(_SCRIPTS_DIR) not in sys.path:
16 sys.path.insert(0, str(_SCRIPTS_DIR))
17
18 from console_encoding import configure_utf8_stdio # noqa: E402
19
20 configure_utf8_stdio()
21
22 if __name__ == "__main__":
23 print(__doc__)
24 print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend fal")
25 raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
26
27 import os
28 import time
29
30 import requests
31
32 from image_backends.backend_common import (
33 MAX_RETRIES,
34 download_image,
35 http_error,
36 is_permanent_error,
37 is_rate_limit_error,
38 normalize_image_size,
39 require_api_key,
40 resolve_output_path,
41 retry_delay,
42 )
43
44
45 VALID_ASPECT_RATIOS = [
46 "1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1",
47 "4:3", "4:5", "5:4", "8:1", "9:16", "16:9", "21:9",
48 ]
49 DEFAULT_ENDPOINT = "https://fal.run"
50 DEFAULT_MODEL = "fal-ai/nano-banana-2"
51 SUPPORTED_MODELS = {DEFAULT_MODEL}
52
53 IMAGE_SIZE_TO_RESOLUTION = {
54 "512px": "0.5K",
55 "1K": "1K",
56 "2K": "2K",
57 "4K": "4K",
58 }
59
60
61 def _resolve_request_options(
62 aspect_ratio: str,
63 image_size: str,
64 model: str,
65 ) -> tuple[str, str]:
66 """Validate request options and return normalized model and resolution values."""
67 resolved_model = model.strip()
68 if resolved_model not in SUPPORTED_MODELS:
69 raise ValueError(
70 f"Unsupported fal model '{model}'. Supported: {sorted(SUPPORTED_MODELS)}"
71 )
72 if aspect_ratio not in VALID_ASPECT_RATIOS:
73 raise ValueError(
74 f"Unsupported aspect ratio '{aspect_ratio}' for fal backend. "
75 f"Supported: {VALID_ASPECT_RATIOS}"
76 )
77 normalized_size = normalize_image_size(image_size)
78 resolution = IMAGE_SIZE_TO_RESOLUTION.get(normalized_size)
79 if not resolution:
80 raise ValueError(
81 f"Unsupported image size '{image_size}' for fal backend. "
82 f"Supported: {list(IMAGE_SIZE_TO_RESOLUTION)}"
83 )
84 return resolved_model, resolution
85
86
87 def _resolve_url(base_url: str, model: str) -> str:
88 """Resolve the full fal endpoint URL for a model."""
89 base = base_url.rstrip("/")
90 if base.endswith(model):
91 return base
92 return f"{base}/{model}"
93
94
95 def _generate_image(api_key: str, prompt: str,
96 aspect_ratio: str = "1:1", image_size: str = "1K",
97 output_dir: str = None, filename: str = None,
98 model: str = DEFAULT_MODEL, base_url: str = DEFAULT_ENDPOINT) -> str:
99 """Generate one image with the fal.ai backend."""
100 model, resolution = _resolve_request_options(aspect_ratio, image_size, model)
101
102 url = _resolve_url(base_url, model)
103 headers = {
104 "Authorization": f"Key {api_key}",
105 "Content-Type": "application/json",
106 }
107 payload = {
108 "prompt": prompt,
109 "aspect_ratio": aspect_ratio,
110 "num_images": 1,
111 "resolution": resolution,
112 }
113
114 print("[fal.ai]")
115 print(f" Model: {model}")
116 print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
117 print(f" Aspect Ratio: {aspect_ratio}")
118 print(f" Resolution: {resolution}")
119 print()
120 print(" [..] Generating...", end="", flush=True)
121 start = time.time()
122 response = requests.post(url, headers=headers, json=payload, timeout=300)
123 elapsed = time.time() - start
124 print(f"\n [DONE] Response received ({elapsed:.1f}s)")
125
126 if response.status_code != 200:
127 raise http_error(response, "fal image generation")
128
129 data = response.json()
130 images = data.get("images") or []
131 image_url = images[0].get("url") if images else None
132 if not image_url:
133 raise RuntimeError(f"fal response missing image URL: {data}")
134
135 path = resolve_output_path(prompt, output_dir, filename, ".png")
136 return download_image(image_url, path)
137
138
139 def generate(prompt: str,
140 aspect_ratio: str = "1:1", image_size: str = "1K",
141 output_dir: str = None, filename: str = None,
142 model: str = None, max_retries: int = MAX_RETRIES) -> str:
143 """Generate an image with retries using the fal.ai backend."""
144 resolved_model = model or os.environ.get("FAL_MODEL") or DEFAULT_MODEL
145 _resolve_request_options(aspect_ratio, image_size, resolved_model)
146 api_key = require_api_key(
147 "FAL_KEY",
148 "FAL_API_KEY",
149 message="No API key found. Set FAL_KEY or FAL_API_KEY in the current environment or a .env file.",
150 )
151 base_url = os.environ.get("FAL_BASE_URL") or DEFAULT_ENDPOINT
152
153 last_error = None
154 for attempt in range(max_retries + 1):
155 try:
156 return _generate_image(
157 api_key=api_key,
158 prompt=prompt,
159 aspect_ratio=aspect_ratio,
160 image_size=image_size,
161 output_dir=output_dir,
162 filename=filename,
163 model=resolved_model,
164 base_url=base_url,
165 )
166 except Exception as exc:
167 last_error = exc
168 if is_permanent_error(exc):
169 raise
170 if attempt >= max_retries:
171 break
172 limited = is_rate_limit_error(exc)
173 delay = retry_delay(attempt, rate_limited=limited)
174 label = "Rate limit hit" if limited else f"Error: {exc}"
175 print(f"\n [WARN] {label}. Retrying in {delay}s...")
176 time.sleep(delay)
177
178 raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
179
179 lines PYTHON