1#!/usr/bin/env python3
2"""Generate face-consistent images using fal.ai Nano Banana 2 edit model.
3 
4Upload a face reference photo and generate new images preserving face identity
5with different expressions, poses, and scenes.
6 
7Auto-fallback: if Nano Banana 2 is down (504/downstream_service_unavailable),
8automatically retries with Nano Banana v1 (Gemini 3 Pro).
9"""
10import argparse
11import json
12import os
13import sys
14import time
15from pathlib import Path
16from urllib.request import urlopen
17 
18MODELS = [
19    "fal-ai/nano-banana-2/edit",
20    "fal-ai/nano-banana/edit",
21]
22 
23 
24def load_fal_key():
25    """Resolve fal.ai credentials in priority order."""
26    for var in ("FAL_KEY", "FAL_API_KEY", "FAL_AI_KEY"):
27        val = os.environ.get(var)
28        if val:
29            os.environ["FAL_KEY"] = val
30            return val
31    for path in ("~/.secrets/fal.env", "~/.fal/key"):
32        p = Path(path).expanduser()
33        if p.exists():
34            for line in p.read_text().splitlines():
35                for prefix in ("FAL_KEY=", "FAL_API_KEY=", "FAL_AI_KEY="):
36                    if line.startswith(prefix):
37                        val = line.split("=", 1)[1].strip()
38                        os.environ["FAL_KEY"] = val
39                        return val
40    raise SystemExit(
41        "Missing fal.ai API key. Set FAL_KEY, FAL_API_KEY, or FAL_AI_KEY env var, "
42        "or pass --api-key."
43    )
44 
45 
46def try_generate(fal_client, model, fal_args):
47    """Try generating with a model. Returns result or raises."""
48    print(f"Generating with {model}...", file=sys.stderr)
49    return fal_client.subscribe(model, arguments=fal_args, with_logs=True)
50 
51 
52def is_retryable(exc):
53    """Check if error is a transient service issue worth retrying/falling back."""
54    msg = str(exc).lower()
55    return any(k in msg for k in [
56        "downstream_service_unavailable",
57        "504",
58        "502",
59        "503",
60        "gateway timeout",
61        "service unavailable",
62    ])
63 
64 
65def main():
66    parser = argparse.ArgumentParser(
67        description="Generate face-consistent images with Nano Banana 2 edit (fal.ai)"
68    )
69    parser.add_argument("--face", required=True, help="Path to face reference image — portrait, rotation grid, or both work")
70    parser.add_argument("--prompt", required=True, help="Scene/expression description prompt")
71    parser.add_argument("--output", required=True, help="Output image path (.png)")
72    parser.add_argument("--aspect-ratio", default="9:16",
73                        help="Aspect ratio: 1:1, 16:9, 9:16, 4:3, 3:4 (default: 9:16)")
74    parser.add_argument("--resolution", default="2K",
75                        help="Resolution: 1K, 2K, 4K (default: 2K)")
76    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility")
77    parser.add_argument("--api-key", help="fal.ai API key (overrides env vars)")
78    parser.add_argument("--model", default=None,
79                        help="Force specific model (skips auto-fallback)")
80    parser.add_argument("--no-fallback", action="store_true",
81                        help="Disable auto-fallback to v1")
82    args = parser.parse_args()
83 
84    if args.api_key:
85        os.environ["FAL_KEY"] = args.api_key
86    else:
87        load_fal_key()
88 
89    import fal_client
90 
91    # Upload face reference if it's a local file
92    face_url = args.face
93    if not face_url.startswith("http"):
94        face_path = Path(face_url).expanduser().resolve()
95        if not face_path.exists():
96            raise SystemExit(f"Face reference not found: {face_path}")
97        print(f"Uploading face reference: {face_path}", file=sys.stderr)
98        face_url = fal_client.upload_file(str(face_path))
99        print(f"Uploaded: {face_url}", file=sys.stderr)
100 
101    image_urls = [face_url]
102 
103    # Build arguments
104    fal_args = {
105        "prompt": args.prompt,
106        "image_urls": image_urls,
107        "aspect_ratio": args.aspect_ratio,
108        "resolution": args.resolution,
109    }
110    if args.seed is not None:
111        fal_args["seed"] = args.seed
112 
113    # Determine which models to try
114    if args.model:
115        models_to_try = [args.model]
116    elif args.no_fallback:
117        models_to_try = [MODELS[0]]
118    else:
119        models_to_try = MODELS
120 
121    # Try each model with retry + fallback
122    result = None
123    last_error = None
124    max_retries = 2
125    for model in models_to_try:
126        for attempt in range(1, max_retries + 1):
127            try:
128                result = try_generate(fal_client, model, fal_args)
129                break
130            except Exception as exc:
131                last_error = exc
132                if "invalid_request" in str(exc).lower():
133                    print(f"  Prompt rejected by {model}. Try simplifying the expression.", file=sys.stderr)
134                    raise
135                if is_retryable(exc):
136                    if attempt < max_retries:
137                        print(f"  {model} attempt {attempt} failed, retrying...", file=sys.stderr)
138                        time.sleep(3)
139                        continue
140                    elif model != models_to_try[-1]:
141                        print(f"  {model} failed {max_retries} times, falling back...", file=sys.stderr)
142                        break
143                    else:
144                        raise
145                else:
146                    raise
147        if result is not None:
148            break
149 
150    if result is None:
151        raise SystemExit(f"All models failed. Last error: {last_error}")
152 
153    # Extract image URL
154    images = result.get("images", [])
155    if not images and "image" in result:
156        images = [result["image"]]
157 
158    if not images:
159        print(json.dumps(result, indent=2), file=sys.stderr)
160        raise SystemExit("No images returned from model")
161 
162    url = images[0].get("url", "")
163    if not url:
164        raise SystemExit("Image has no URL")
165 
166    # Download
167    output_path = Path(args.output).expanduser().resolve()
168    output_path.parent.mkdir(parents=True, exist_ok=True)
169    with urlopen(url, timeout=120) as resp:
170        output_path.write_bytes(resp.read())
171 
172    print(json.dumps({
173        "ok": True,
174        "output": str(output_path),
175        "url": url,
176        "model": model,
177        "prompt": args.prompt,
178    }))
179 
180 
181if __name__ == "__main__":
182    main()
183