返回 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)
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_rate_limit_error,
37 require_api_key,
38 resolve_output_path,
39 retry_delay,
40 )
41
42
43 VALID_ASPECT_RATIOS = ["1:1", "16:9", "9:16", "3:4", "4:3"]
44 DEFAULT_ENDPOINT = "https://fal.run"
45 DEFAULT_MODEL = "fal-ai/imagen3/fast"
46
47
48 def _resolve_url(base_url: str, model: str) -> str:
49 """Resolve the full fal endpoint URL for a model."""
50 base = base_url.rstrip("/")
51 if base.endswith(model):
52 return base
53 return f"{base}/{model}"
54
55
56 def _generate_image(api_key: str, prompt: str,
57 aspect_ratio: str = "1:1", image_size: str = "1K",
58 output_dir: str = None, filename: str = None,
59 model: str = DEFAULT_MODEL, base_url: str = DEFAULT_ENDPOINT) -> str:
60 """Generate one image with the fal.ai backend."""
61 del image_size
62
63 if aspect_ratio not in VALID_ASPECT_RATIOS:
64 raise ValueError(
65 f"Unsupported aspect ratio '{aspect_ratio}' for fal backend. "
66 f"Supported: {VALID_ASPECT_RATIOS}"
67 )
68
69 url = _resolve_url(base_url, model)
70 headers = {
71 "Authorization": f"Key {api_key}",
72 "Content-Type": "application/json",
73 }
74 payload = {
75 "prompt": prompt,
76 "aspect_ratio": aspect_ratio,
77 "num_images": 1,
78 }
79
80 print("[fal.ai]")
81 print(f" Model: {model}")
82 print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
83 print(f" Aspect Ratio: {aspect_ratio}")
84 print()
85 print(" [..] Generating...", end="", flush=True)
86 start = time.time()
87 response = requests.post(url, headers=headers, json=payload, timeout=300)
88 elapsed = time.time() - start
89 print(f"\n [DONE] Response received ({elapsed:.1f}s)")
90
91 if response.status_code != 200:
92 raise http_error(response, "fal image generation")
93
94 data = response.json()
95 images = data.get("images") or []
96 image_url = images[0].get("url") if images else None
97 if not image_url:
98 raise RuntimeError(f"fal response missing image URL: {data}")
99
100 path = resolve_output_path(prompt, output_dir, filename, ".png")
101 return download_image(image_url, path)
102
103
104 def generate(prompt: str,
105 aspect_ratio: str = "1:1", image_size: str = "1K",
106 output_dir: str = None, filename: str = None,
107 model: str = None, max_retries: int = MAX_RETRIES) -> str:
108 """Generate an image with retries using the fal.ai backend."""
109 api_key = require_api_key(
110 "FAL_KEY",
111 "FAL_API_KEY",
112 message="No API key found. Set FAL_KEY or FAL_API_KEY in the current environment or a .env file.",
113 )
114 base_url = os.environ.get("FAL_BASE_URL") or DEFAULT_ENDPOINT
115 resolved_model = model or os.environ.get("FAL_MODEL") or DEFAULT_MODEL
116
117 last_error = None
118 for attempt in range(max_retries + 1):
119 try:
120 return _generate_image(
121 api_key=api_key,
122 prompt=prompt,
123 aspect_ratio=aspect_ratio,
124 image_size=image_size,
125 output_dir=output_dir,
126 filename=filename,
127 model=resolved_model,
128 base_url=base_url,
129 )
130 except Exception as exc:
131 last_error = exc
132 if attempt >= max_retries:
133 break
134 limited = is_rate_limit_error(exc)
135 delay = retry_delay(attempt, rate_limited=limited)
136 label = "Rate limit hit" if limited else f"Error: {exc}"
137 print(f"\n [WARN] {label}. Retrying in {delay}s...")
138 time.sleep(delay)
139
140 raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
141
141 lines PYTHON