google-flow-mcp 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,33 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [main, "claude/**"]
6
+ pull_request:
7
+ branches: [main]
8
+
9
+ jobs:
10
+ test:
11
+ name: Test (Python ${{ matrix.python-version }})
12
+ runs-on: ubuntu-latest
13
+ strategy:
14
+ fail-fast: false
15
+ matrix:
16
+ python-version: ["3.10", "3.11", "3.12"]
17
+
18
+ steps:
19
+ - uses: actions/checkout@v4
20
+
21
+ - name: Set up Python ${{ matrix.python-version }}
22
+ uses: actions/setup-python@v5
23
+ with:
24
+ python-version: ${{ matrix.python-version }}
25
+
26
+ - name: Install uv
27
+ uses: astral-sh/setup-uv@v4
28
+
29
+ - name: Install package with dev dependencies
30
+ run: uv pip install --system -e ".[dev]"
31
+
32
+ - name: Run tests
33
+ run: pytest -v --tb=short tests/
@@ -0,0 +1,15 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ dist/
5
+ build/
6
+ .venv/
7
+ .env
8
+ *.egg
9
+ .pytest_cache/
10
+ .ruff_cache/
11
+ google_flow_outputs/
12
+ *.mp4
13
+ *.png
14
+ *.jpg
15
+ *.jpeg
@@ -0,0 +1,139 @@
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship made available under
36
+ the License, as indicated by a copyright notice that is included in
37
+ or attached to the work (an example is provided in the Appendix below).
38
+
39
+ "Derivative Works" shall mean any work, whether in Source or Object
40
+ form, that is based on (or derived from) the Work and for which the
41
+ editorial revisions, annotations, elaborations, or other modifications
42
+ represent, as a whole, an original work of authorship. For the purposes
43
+ of this License, Derivative Works shall not include works that remain
44
+ separable from, or merely link (or bind by name) to the interfaces of,
45
+ the Work and Derivative Works thereof.
46
+
47
+ "Contribution" shall mean, as submitted to the Licensor for inclusion
48
+ in the Work by the copyright owner or by an individual or Legal Entity
49
+ authorized to submit on behalf of the copyright owner.
50
+
51
+ "Contributor" shall mean Licensor and any Legal Entity on behalf of
52
+ whom a Contribution has been received by the Licensor and included
53
+ within the Work.
54
+
55
+ 2. Grant of Copyright License. Subject to the terms and conditions of
56
+ this License, each Contributor hereby grants to You a perpetual,
57
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
58
+ copyright license to reproduce, prepare Derivative Works of,
59
+ publicly display, publicly perform, sublicense, and distribute the
60
+ Work and such Derivative Works in Source or Object form.
61
+
62
+ 3. Grant of Patent License. Subject to the terms and conditions of
63
+ this License, each Contributor hereby grants to You a perpetual,
64
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
65
+ (except as stated in this section) patent license to make, have made,
66
+ use, offer to sell, sell, import, and otherwise transfer the Work,
67
+ where such license applies only to those patent claims licensable
68
+ by such Contributor that are necessarily infringed by their
69
+ Contribution(s) alone or by the combination of their Contribution(s)
70
+ with the Work to which such Contribution(s) was submitted. If You
71
+ institute patent litigation against any entity (including a cross-claim
72
+ or counterclaim in a lawsuit) alleging that the Work or any Contribution
73
+ embodied within the Work constitutes direct or contributory patent
74
+ infringement, then any patent licenses granted to You under this License
75
+ for that Work shall terminate as of the date such litigation is filed.
76
+
77
+ 4. Redistribution. You may reproduce and distribute copies of the
78
+ Work or Derivative Works thereof in any medium, with or without
79
+ modifications, and in Source or Object form, provided that You
80
+ meet the following conditions:
81
+
82
+ (a) You must give any other recipients of the Work or Derivative Works
83
+ a copy of this License; and
84
+
85
+ (b) You must cause any modified files to carry prominent notices
86
+ stating that You changed the files; and
87
+
88
+ (c) You must retain, in the Source form of any Derivative Works
89
+ that You distribute, all copyright, patent, trademark, and
90
+ attribution notices from the Source form of the Work; and
91
+
92
+ (d) If the Work includes a "NOTICE" text file, you must include a
93
+ readable copy of the attribution notices contained within such
94
+ NOTICE file.
95
+
96
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
97
+ any Contribution intentionally submitted for inclusion in the Work
98
+ by You to the Licensor shall be under the terms and conditions of
99
+ this License, without any additional terms or conditions.
100
+
101
+ 6. Trademarks. This License does not grant permission to use the trade
102
+ names, trademarks, service marks, or product names of the Licensor.
103
+
104
+ 7. Disclaimer of Warranty. Unless required by applicable law or agreed
105
+ to in writing, Licensor provides the Work on an "AS IS" BASIS,
106
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
107
+ implied. You are solely responsible for determining the
108
+ appropriateness of using or reproducing the Work.
109
+
110
+ 8. Limitation of Liability. In no event and under no legal theory,
111
+ whether in tort (including negligence), contract, or otherwise,
112
+ unless required by applicable law (such as deliberate and grossly
113
+ negligent acts) shall any Contributor be liable to You for damages,
114
+ including any direct, indirect, special, incidental, or exemplary
115
+ damages of any character arising as a result of this License or out
116
+ of the use or inability to use the Work.
117
+
118
+ 9. Accepting Warranty or Additional Liability. While redistributing
119
+ the Work or Derivative Works thereof, You may choose to offer, and
120
+ charge a fee for, acceptance of support, warranty, indemnity, or
121
+ other liability obligations and/or rights consistent with this
122
+ License. However, in accepting such obligations, You may offer such
123
+ obligations only on Your own behalf and on Your sole responsibility.
124
+
125
+ END OF TERMS AND CONDITIONS
126
+
127
+ Copyright 2025 Joshua Daniel
128
+
129
+ Licensed under the Apache License, Version 2.0 (the "License");
130
+ you may not use this file except in compliance with the License.
131
+ You may obtain a copy of the License at
132
+
133
+ http://www.apache.org/licenses/LICENSE-2.0
134
+
135
+ Unless required by applicable law or agreed to in writing, software
136
+ distributed under the License is distributed on an "AS IS" BASIS,
137
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
138
+ See the License for the specific language governing permissions and
139
+ limitations under the License.
@@ -0,0 +1,134 @@
1
+ Metadata-Version: 2.4
2
+ Name: google-flow-mcp
3
+ Version: 0.1.0
4
+ Summary: MCP server giving Claude Desktop access to Google AI image and video generation
5
+ Project-URL: Homepage, https://github.com/joshuadaniel-8090/google-flow-mcp
6
+ Project-URL: Issues, https://github.com/joshuadaniel-8090/google-flow-mcp/issues
7
+ License: Apache-2.0
8
+ License-File: LICENSE
9
+ Keywords: claude,google-ai,image-generation,imagen,mcp,veo,video-generation
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: License :: OSI Approved :: Apache Software License
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Programming Language :: Python :: 3.10
15
+ Classifier: Programming Language :: Python :: 3.11
16
+ Classifier: Programming Language :: Python :: 3.12
17
+ Requires-Python: >=3.10
18
+ Requires-Dist: google-genai>=1.0.0
19
+ Requires-Dist: mcp[cli]>=1.0.0
20
+ Requires-Dist: pillow>=10.0.0
21
+ Provides-Extra: dev
22
+ Requires-Dist: pytest-asyncio>=0.23.0; extra == 'dev'
23
+ Requires-Dist: pytest-mock>=3.12.0; extra == 'dev'
24
+ Requires-Dist: pytest>=8.0.0; extra == 'dev'
25
+ Description-Content-Type: text/markdown
26
+
27
+ # Google Flow MCP
28
+
29
+ A Python MCP server that gives **Claude Desktop** direct access to Google AI's latest image and video generation models.
30
+
31
+ | Capability | Model | Tier |
32
+ |---|---|---|
33
+ | High-quality image generation & editing | **Nano Banana Pro** (`gemini-3-pro-image`) | Free |
34
+ | Fast image generation | **Nano Banana 2** (`gemini-3.1-flash-image`) | Free |
35
+ | Cinematic video with native audio | **Veo 3.1** (`veo-3.1-generate-preview`) | Paid |
36
+
37
+ ## Tools
38
+
39
+ | Tool | What it does |
40
+ |---|---|
41
+ | `flow_generate_image` | Text → 1–4 images, up to 4K, choice of model |
42
+ | `flow_edit_image` | Edit an image with natural language (inpaint, outpaint, bg-swap) |
43
+ | `flow_generate_image_with_references` | Generate guided by up to 14 reference images |
44
+ | `flow_generate_video` | Text → cinematic video, optional anchor frame |
45
+ | `flow_extend_video` | Extend an existing Veo clip |
46
+ | `flow_image_to_video` | Full pipeline: Nano Banana Pro image → Veo 3.1 video |
47
+
48
+ ## Prerequisites
49
+
50
+ - Python 3.10 or newer
51
+ - [`uv`](https://docs.astral.sh/uv/getting-started/installation/) (recommended) or `pip`
52
+ - A Google AI Studio API key — get one free at [aistudio.google.com/apikey](https://aistudio.google.com/apikey)
53
+
54
+ ## Installation
55
+
56
+ ```bash
57
+ # Using uv (recommended — no virtual env setup needed)
58
+ uvx google-flow-mcp
59
+
60
+ # Or install with pip
61
+ pip install google-flow-mcp
62
+ ```
63
+
64
+ ## Claude Desktop Configuration
65
+
66
+ Add the following to your Claude Desktop config file:
67
+
68
+ ```json
69
+ {
70
+ "mcpServers": {
71
+ "google-flow": {
72
+ "command": "uvx",
73
+ "args": ["google-flow-mcp"],
74
+ "env": {
75
+ "GOOGLE_API_KEY": "YOUR_GOOGLE_AI_STUDIO_KEY"
76
+ }
77
+ }
78
+ }
79
+ }
80
+ ```
81
+
82
+ ### Config file locations
83
+
84
+ | Platform | Path |
85
+ |---|---|
86
+ | **Windows Store** | `%LOCALAPPDATA%\Packages\Claude_pzs8sxrjxfjjc\LocalCache\Roaming\Claude\claude_desktop_config.json` |
87
+ | **Windows Direct** | `%APPDATA%\Claude\claude_desktop_config.json` |
88
+ | **macOS** | `~/Library/Application Support/Claude/claude_desktop_config.json` |
89
+ | **Linux** | `~/.config/Claude/claude_desktop_config.json` |
90
+
91
+ After editing the config, **restart Claude Desktop**.
92
+
93
+ ## Output Directory
94
+
95
+ Generated files are saved to `~/google_flow_outputs/` by default.
96
+
97
+ Override with the `FLOW_OUTPUT_DIR` environment variable:
98
+
99
+ ```json
100
+ "env": {
101
+ "GOOGLE_API_KEY": "YOUR_KEY",
102
+ "FLOW_OUTPUT_DIR": "/Users/you/Pictures/ai-outputs"
103
+ }
104
+ ```
105
+
106
+ ## Usage Examples
107
+
108
+ Once connected, ask Claude naturally:
109
+
110
+ - *"Generate a photo-realistic image of a neon-lit Tokyo alley at night"*
111
+ - *"Edit this image to remove the background and replace it with a forest"*
112
+ - *"Generate a video of a rocket launching from a desert at dusk with dramatic audio"*
113
+ - *"Create an anchor image of ocean waves, then animate it into a video"*
114
+
115
+ ## Important Notes
116
+
117
+ - **Image generation** (Nano Banana Pro / Nano Banana 2) works on the **free tier**.
118
+ - **Video generation** (Veo 3.1) requires a **paid Google AI API plan**.
119
+ - All Veo videos are **SynthID-watermarked** by Google.
120
+ - Generated videos are stored on Google's servers for **2 days** after creation.
121
+ - Video generation typically takes **1–4 minutes** — Claude will wait automatically.
122
+
123
+ ## Development
124
+
125
+ ```bash
126
+ git clone https://github.com/joshuadaniel-8090/google-flow-mcp
127
+ cd google-flow-mcp
128
+ pip install -e ".[dev]"
129
+ pytest -v
130
+ ```
131
+
132
+ ## License
133
+
134
+ Apache 2.0 — see [LICENSE](LICENSE).
@@ -0,0 +1,108 @@
1
+ # Google Flow MCP
2
+
3
+ A Python MCP server that gives **Claude Desktop** direct access to Google AI's latest image and video generation models.
4
+
5
+ | Capability | Model | Tier |
6
+ |---|---|---|
7
+ | High-quality image generation & editing | **Nano Banana Pro** (`gemini-3-pro-image`) | Free |
8
+ | Fast image generation | **Nano Banana 2** (`gemini-3.1-flash-image`) | Free |
9
+ | Cinematic video with native audio | **Veo 3.1** (`veo-3.1-generate-preview`) | Paid |
10
+
11
+ ## Tools
12
+
13
+ | Tool | What it does |
14
+ |---|---|
15
+ | `flow_generate_image` | Text → 1–4 images, up to 4K, choice of model |
16
+ | `flow_edit_image` | Edit an image with natural language (inpaint, outpaint, bg-swap) |
17
+ | `flow_generate_image_with_references` | Generate guided by up to 14 reference images |
18
+ | `flow_generate_video` | Text → cinematic video, optional anchor frame |
19
+ | `flow_extend_video` | Extend an existing Veo clip |
20
+ | `flow_image_to_video` | Full pipeline: Nano Banana Pro image → Veo 3.1 video |
21
+
22
+ ## Prerequisites
23
+
24
+ - Python 3.10 or newer
25
+ - [`uv`](https://docs.astral.sh/uv/getting-started/installation/) (recommended) or `pip`
26
+ - A Google AI Studio API key — get one free at [aistudio.google.com/apikey](https://aistudio.google.com/apikey)
27
+
28
+ ## Installation
29
+
30
+ ```bash
31
+ # Using uv (recommended — no virtual env setup needed)
32
+ uvx google-flow-mcp
33
+
34
+ # Or install with pip
35
+ pip install google-flow-mcp
36
+ ```
37
+
38
+ ## Claude Desktop Configuration
39
+
40
+ Add the following to your Claude Desktop config file:
41
+
42
+ ```json
43
+ {
44
+ "mcpServers": {
45
+ "google-flow": {
46
+ "command": "uvx",
47
+ "args": ["google-flow-mcp"],
48
+ "env": {
49
+ "GOOGLE_API_KEY": "YOUR_GOOGLE_AI_STUDIO_KEY"
50
+ }
51
+ }
52
+ }
53
+ }
54
+ ```
55
+
56
+ ### Config file locations
57
+
58
+ | Platform | Path |
59
+ |---|---|
60
+ | **Windows Store** | `%LOCALAPPDATA%\Packages\Claude_pzs8sxrjxfjjc\LocalCache\Roaming\Claude\claude_desktop_config.json` |
61
+ | **Windows Direct** | `%APPDATA%\Claude\claude_desktop_config.json` |
62
+ | **macOS** | `~/Library/Application Support/Claude/claude_desktop_config.json` |
63
+ | **Linux** | `~/.config/Claude/claude_desktop_config.json` |
64
+
65
+ After editing the config, **restart Claude Desktop**.
66
+
67
+ ## Output Directory
68
+
69
+ Generated files are saved to `~/google_flow_outputs/` by default.
70
+
71
+ Override with the `FLOW_OUTPUT_DIR` environment variable:
72
+
73
+ ```json
74
+ "env": {
75
+ "GOOGLE_API_KEY": "YOUR_KEY",
76
+ "FLOW_OUTPUT_DIR": "/Users/you/Pictures/ai-outputs"
77
+ }
78
+ ```
79
+
80
+ ## Usage Examples
81
+
82
+ Once connected, ask Claude naturally:
83
+
84
+ - *"Generate a photo-realistic image of a neon-lit Tokyo alley at night"*
85
+ - *"Edit this image to remove the background and replace it with a forest"*
86
+ - *"Generate a video of a rocket launching from a desert at dusk with dramatic audio"*
87
+ - *"Create an anchor image of ocean waves, then animate it into a video"*
88
+
89
+ ## Important Notes
90
+
91
+ - **Image generation** (Nano Banana Pro / Nano Banana 2) works on the **free tier**.
92
+ - **Video generation** (Veo 3.1) requires a **paid Google AI API plan**.
93
+ - All Veo videos are **SynthID-watermarked** by Google.
94
+ - Generated videos are stored on Google's servers for **2 days** after creation.
95
+ - Video generation typically takes **1–4 minutes** — Claude will wait automatically.
96
+
97
+ ## Development
98
+
99
+ ```bash
100
+ git clone https://github.com/joshuadaniel-8090/google-flow-mcp
101
+ cd google-flow-mcp
102
+ pip install -e ".[dev]"
103
+ pytest -v
104
+ ```
105
+
106
+ ## License
107
+
108
+ Apache 2.0 — see [LICENSE](LICENSE).
@@ -0,0 +1,47 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "google-flow-mcp"
7
+ version = "0.1.0"
8
+ description = "MCP server giving Claude Desktop access to Google AI image and video generation"
9
+ readme = "README.md"
10
+ license = { text = "Apache-2.0" }
11
+ requires-python = ">=3.10"
12
+ keywords = ["mcp", "google-ai", "imagen", "veo", "claude", "image-generation", "video-generation"]
13
+ classifiers = [
14
+ "Development Status :: 3 - Alpha",
15
+ "Intended Audience :: Developers",
16
+ "License :: OSI Approved :: Apache Software License",
17
+ "Programming Language :: Python :: 3",
18
+ "Programming Language :: Python :: 3.10",
19
+ "Programming Language :: Python :: 3.11",
20
+ "Programming Language :: Python :: 3.12",
21
+ ]
22
+ dependencies = [
23
+ "mcp[cli]>=1.0.0",
24
+ "google-genai>=1.0.0",
25
+ "Pillow>=10.0.0",
26
+ ]
27
+
28
+ [project.optional-dependencies]
29
+ dev = [
30
+ "pytest>=8.0.0",
31
+ "pytest-asyncio>=0.23.0",
32
+ "pytest-mock>=3.12.0",
33
+ ]
34
+
35
+ [project.scripts]
36
+ google-flow-mcp = "google_flow_mcp.server:main"
37
+
38
+ [project.urls]
39
+ Homepage = "https://github.com/joshuadaniel-8090/google-flow-mcp"
40
+ Issues = "https://github.com/joshuadaniel-8090/google-flow-mcp/issues"
41
+
42
+ [tool.hatch.build.targets.wheel]
43
+ packages = ["src/google_flow_mcp"]
44
+
45
+ [tool.pytest.ini_options]
46
+ asyncio_mode = "auto"
47
+ testpaths = ["tests"]
@@ -0,0 +1,3 @@
1
+ """Google Flow MCP — Google AI image and video generation for Claude Desktop."""
2
+
3
+ __version__ = "0.1.0"
@@ -0,0 +1,551 @@
1
+ """Google Flow MCP server — 6 tools for Google AI image and video generation."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import asyncio
6
+ import io
7
+ import mimetypes
8
+ import os
9
+ import time
10
+ from datetime import datetime
11
+ from pathlib import Path
12
+ from typing import Optional
13
+
14
+ from google import genai
15
+ from google.genai import types
16
+ from mcp.server.fastmcp import FastMCP
17
+ from PIL import Image as PILImage
18
+
19
+ # ── Constants ────────────────────────────────────────────────────────────────
20
+
21
+ MODEL_IMAGE_PRO = "gemini-3-pro-image" # Nano Banana Pro
22
+ MODEL_IMAGE_FLASH = "gemini-3.1-flash-image" # Nano Banana 2
23
+ MODEL_VIDEO = "veo-3.1-generate-preview" # Veo 3.1
24
+
25
+ VIDEO_POLL_INTERVAL = 10 # seconds between Veo status polls
26
+ VIDEO_POLL_TIMEOUT = 600 # 10-minute max wait for video generation
27
+
28
+ # ── FastMCP Instance ─────────────────────────────────────────────────────────
29
+
30
+ mcp = FastMCP(
31
+ name="google-flow-mcp",
32
+ instructions=(
33
+ "Provides 6 tools for Google AI image and video generation. "
34
+ "Image tools use Nano Banana Pro (gemini-3-pro-image, high quality) or "
35
+ "Nano Banana 2 (gemini-3.1-flash-image, fast). "
36
+ "Video tools use Veo 3.1 (veo-3.1-generate-preview) with native audio. "
37
+ "Image generation works on the free tier. "
38
+ "Video generation requires a paid Google AI API plan. "
39
+ "Outputs are saved to ~/google_flow_outputs/ by default (override with FLOW_OUTPUT_DIR)."
40
+ ),
41
+ )
42
+
43
+ # ── Client ────────────────────────────────────────────────────────────────────
44
+
45
+ _client: genai.Client | None = None
46
+
47
+
48
+ def _get_client() -> genai.Client:
49
+ api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
50
+ if not api_key:
51
+ raise EnvironmentError(
52
+ "GOOGLE_API_KEY environment variable is required. "
53
+ "Get a free key at https://aistudio.google.com/apikey"
54
+ )
55
+ return genai.Client(api_key=api_key)
56
+
57
+
58
+ def get_client() -> genai.Client:
59
+ global _client
60
+ if _client is None:
61
+ _client = _get_client()
62
+ return _client
63
+
64
+
65
+ # ── File Helpers ──────────────────────────────────────────────────────────────
66
+
67
+
68
+ def _output_dir() -> Path:
69
+ base = os.environ.get("FLOW_OUTPUT_DIR") or str(Path.home() / "google_flow_outputs")
70
+ path = Path(base)
71
+ path.mkdir(parents=True, exist_ok=True)
72
+ return path
73
+
74
+
75
+ def _timestamped(prefix: str, ext: str, out: Optional[Path] = None) -> Path:
76
+ ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
77
+ return (out or _output_dir()) / f"{prefix}_{ts}.{ext}"
78
+
79
+
80
+ def _save_image(image_obj, path: Path) -> None:
81
+ """Save a generated image to disk, handling multiple SDK response shapes."""
82
+ if hasattr(image_obj, "save"):
83
+ image_obj.save(str(path))
84
+ elif getattr(image_obj, "image_bytes", None):
85
+ PILImage.open(io.BytesIO(image_obj.image_bytes)).save(str(path))
86
+ elif getattr(image_obj, "_image_bytes", None):
87
+ PILImage.open(io.BytesIO(image_obj._image_bytes)).save(str(path))
88
+ else:
89
+ raise RuntimeError(f"Cannot save image of type {type(image_obj).__name__}")
90
+
91
+
92
+ def _load_image_bytes(path: str) -> bytes:
93
+ with open(path, "rb") as f:
94
+ return f.read()
95
+
96
+
97
+ def _mime(path: str) -> str:
98
+ mime, _ = mimetypes.guess_type(path)
99
+ return mime or "image/png"
100
+
101
+
102
+ def _resolve_image_model(model: str) -> str:
103
+ return MODEL_IMAGE_FLASH if model.lower() in ("flash", "2") else MODEL_IMAGE_PRO
104
+
105
+
106
+ async def _poll_operation(client: genai.Client, operation) -> object:
107
+ """Poll a Veo long-running operation until done or timeout."""
108
+ deadline = time.monotonic() + VIDEO_POLL_TIMEOUT
109
+ while not operation.done:
110
+ if time.monotonic() > deadline:
111
+ raise TimeoutError(
112
+ f"Video generation timed out after {VIDEO_POLL_TIMEOUT // 60} minutes. "
113
+ "Try a shorter duration or check the Google AI console."
114
+ )
115
+ await asyncio.sleep(VIDEO_POLL_INTERVAL)
116
+ operation = await client.aio.operations.get(operation)
117
+ if getattr(operation, "error", None) and operation.error.message:
118
+ raise RuntimeError(f"Video generation failed: {operation.error.message}")
119
+ return operation
120
+
121
+
122
+ def _collect_video_results(operation, out: Path, ts: str) -> tuple[list[str], list[str]]:
123
+ """Extract saved paths and/or URIs from a completed video operation."""
124
+ saved, uris = [], []
125
+ videos = getattr(operation.response, "generated_videos", [])
126
+ for i, gen_video in enumerate(videos):
127
+ vid = gen_video.video
128
+ vid_bytes = getattr(vid, "video_bytes", None)
129
+ uri = getattr(vid, "uri", None)
130
+ if vid_bytes:
131
+ path = out / f"flow_video_{ts}_{i}.mp4"
132
+ path.write_bytes(vid_bytes)
133
+ saved.append(str(path))
134
+ elif uri:
135
+ uris.append(uri)
136
+ return saved, uris
137
+
138
+
139
+ def _format_video_result(saved: list[str], uris: list[str]) -> str:
140
+ parts = []
141
+ if saved:
142
+ parts.append("Saved videos:\n" + "\n".join(saved))
143
+ if uris:
144
+ parts.append(
145
+ "Video URI(s) (available for 2 days on Google servers):\n" + "\n".join(uris)
146
+ )
147
+ parts.append("Note: All Veo videos are SynthID-watermarked by Google.")
148
+ return "\n\n".join(parts)
149
+
150
+
151
+ # ── Tool 1: Text → Image ─────────────────────────────────────────────────────
152
+
153
+
154
+ @mcp.tool()
155
+ async def flow_generate_image(
156
+ prompt: str,
157
+ model: str = "pro",
158
+ number_of_images: int = 1,
159
+ aspect_ratio: str = "1:1",
160
+ negative_prompt: Optional[str] = None,
161
+ output_dir: Optional[str] = None,
162
+ ) -> str:
163
+ """
164
+ Generate 1–4 images from a text prompt using Google AI.
165
+
166
+ Args:
167
+ prompt: Detailed description of the image(s) to create.
168
+ model: 'pro' for Nano Banana Pro (high quality) or 'flash' for Nano Banana 2 (fast).
169
+ number_of_images: How many images to generate (1–4).
170
+ aspect_ratio: '1:1', '3:4', '4:3', '9:16', or '16:9'.
171
+ negative_prompt: What to exclude from the generated images.
172
+ output_dir: Override the output directory (default: ~/google_flow_outputs).
173
+
174
+ Returns:
175
+ Paths to the saved image files.
176
+ """
177
+ client = get_client()
178
+ model_id = _resolve_image_model(model)
179
+ n = max(1, min(4, number_of_images))
180
+ out = Path(output_dir) if output_dir else _output_dir()
181
+ out.mkdir(parents=True, exist_ok=True)
182
+
183
+ response = await client.aio.models.generate_images(
184
+ model=model_id,
185
+ prompt=prompt,
186
+ config=types.GenerateImagesConfig(
187
+ number_of_images=n,
188
+ aspect_ratio=aspect_ratio,
189
+ negative_prompt=negative_prompt,
190
+ ),
191
+ )
192
+
193
+ ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
194
+ paths = []
195
+ for i, gen_img in enumerate(response.generated_images):
196
+ path = out / f"flow_image_{ts}_{i}.png"
197
+ _save_image(gen_img.image, path)
198
+ paths.append(str(path))
199
+
200
+ return f"Generated {len(paths)} image(s):\n" + "\n".join(paths)
201
+
202
+
203
+ # ── Tool 2: Edit Image ────────────────────────────────────────────────────────
204
+
205
+
206
+ @mcp.tool()
207
+ async def flow_edit_image(
208
+ image_path: str,
209
+ prompt: str,
210
+ edit_mode: str = "inpaint_insertion",
211
+ model: str = "pro",
212
+ mask_mode: str = "MASK_MODE_BACKGROUND",
213
+ output_dir: Optional[str] = None,
214
+ ) -> str:
215
+ """
216
+ Edit an existing image using a natural language instruction.
217
+
218
+ Args:
219
+ image_path: Absolute path to the source image file.
220
+ prompt: Natural language description of the edit to apply.
221
+ edit_mode: 'inpaint_insertion' (add), 'inpaint_removal' (remove),
222
+ 'outpaint' (expand canvas), 'bgswap' (replace background).
223
+ model: 'pro' for Nano Banana Pro or 'flash' for Nano Banana 2.
224
+ mask_mode: Auto-mask strategy — 'MASK_MODE_BACKGROUND', 'MASK_MODE_FOREGROUND',
225
+ or 'MASK_MODE_SEMANTIC'.
226
+ output_dir: Override the output directory (default: ~/google_flow_outputs).
227
+
228
+ Returns:
229
+ Path to the edited image file.
230
+ """
231
+ client = get_client()
232
+ model_id = _resolve_image_model(model)
233
+ out = Path(output_dir) if output_dir else _output_dir()
234
+ out.mkdir(parents=True, exist_ok=True)
235
+
236
+ _edit_mode_map = {
237
+ "inpaint_insertion": "EDIT_MODE_INPAINT_INSERTION",
238
+ "inpaint_removal": "EDIT_MODE_INPAINT_REMOVAL",
239
+ "outpaint": "EDIT_MODE_OUTPAINT",
240
+ "bgswap": "EDIT_MODE_BGSWAP",
241
+ }
242
+ sdk_edit_mode = _edit_mode_map.get(edit_mode, edit_mode)
243
+
244
+ raw_ref = types.RawReferenceImage(
245
+ reference_id=0,
246
+ reference_image=types.Image(
247
+ image_bytes=_load_image_bytes(image_path),
248
+ mime_type=_mime(image_path),
249
+ ),
250
+ )
251
+ mask_ref = types.MaskReferenceImage(
252
+ reference_id=1,
253
+ config=types.MaskReferenceConfig(mask_mode=mask_mode),
254
+ )
255
+
256
+ response = await client.aio.models.edit_image(
257
+ model=model_id,
258
+ prompt=prompt,
259
+ reference_images=[raw_ref, mask_ref],
260
+ config=types.EditImageConfig(
261
+ edit_mode=sdk_edit_mode,
262
+ number_of_images=1,
263
+ ),
264
+ )
265
+
266
+ ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
267
+ path = out / f"flow_edit_{ts}.png"
268
+ _save_image(response.generated_images[0].image, path)
269
+ return f"Edited image saved to:\n{path}"
270
+
271
+
272
+ # ── Tool 3: Generate with Reference Images ────────────────────────────────────
273
+
274
+
275
+ @mcp.tool()
276
+ async def flow_generate_image_with_references(
277
+ prompt: str,
278
+ reference_image_paths: list[str],
279
+ model: str = "pro",
280
+ aspect_ratio: str = "1:1",
281
+ number_of_images: int = 1,
282
+ output_dir: Optional[str] = None,
283
+ ) -> str:
284
+ """
285
+ Generate an image guided by up to 14 reference images.
286
+
287
+ The model uses the reference images for style, composition, and subject guidance
288
+ while following the text prompt.
289
+
290
+ Args:
291
+ prompt: Description of the image to generate.
292
+ reference_image_paths: List of paths to reference image files (max 14).
293
+ model: 'pro' for Nano Banana Pro or 'flash' for Nano Banana 2.
294
+ aspect_ratio: '1:1', '3:4', '4:3', '9:16', or '16:9'.
295
+ number_of_images: How many images to generate (1–4).
296
+ output_dir: Override the output directory (default: ~/google_flow_outputs).
297
+
298
+ Returns:
299
+ Paths to the generated image files.
300
+ """
301
+ if len(reference_image_paths) > 14:
302
+ raise ValueError(
303
+ f"Maximum 14 reference images allowed; got {len(reference_image_paths)}."
304
+ )
305
+ if not reference_image_paths:
306
+ raise ValueError("At least one reference image path is required.")
307
+
308
+ client = get_client()
309
+ model_id = _resolve_image_model(model)
310
+ n = max(1, min(4, number_of_images))
311
+ out = Path(output_dir) if output_dir else _output_dir()
312
+ out.mkdir(parents=True, exist_ok=True)
313
+
314
+ reference_images = [
315
+ types.RawReferenceImage(
316
+ reference_id=i,
317
+ reference_image=types.Image(
318
+ image_bytes=_load_image_bytes(ref_path),
319
+ mime_type=_mime(ref_path),
320
+ ),
321
+ )
322
+ for i, ref_path in enumerate(reference_image_paths)
323
+ ]
324
+
325
+ response = await client.aio.models.generate_images(
326
+ model=model_id,
327
+ prompt=prompt,
328
+ reference_images=reference_images,
329
+ config=types.GenerateImagesConfig(
330
+ number_of_images=n,
331
+ aspect_ratio=aspect_ratio,
332
+ ),
333
+ )
334
+
335
+ ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
336
+ paths = []
337
+ for i, gen_img in enumerate(response.generated_images):
338
+ path = out / f"flow_ref_image_{ts}_{i}.png"
339
+ _save_image(gen_img.image, path)
340
+ paths.append(str(path))
341
+
342
+ return (
343
+ f"Generated {len(paths)} image(s) using {len(reference_image_paths)} reference(s):\n"
344
+ + "\n".join(paths)
345
+ )
346
+
347
+
348
+ # ── Tool 4: Text → Video ─────────────────────────────────────────────────────
349
+
350
+
351
+ @mcp.tool()
352
+ async def flow_generate_video(
353
+ prompt: str,
354
+ duration_seconds: int = 8,
355
+ aspect_ratio: str = "16:9",
356
+ negative_prompt: Optional[str] = None,
357
+ anchor_frame_path: Optional[str] = None,
358
+ enhance_prompt: bool = True,
359
+ output_dir: Optional[str] = None,
360
+ ) -> str:
361
+ """
362
+ Generate a cinematic video with native audio from a text prompt using Veo 3.1.
363
+
364
+ Requires a paid Google AI API plan. Generated videos are SynthID-watermarked
365
+ and stored on Google servers for 2 days.
366
+
367
+ Args:
368
+ prompt: Cinematic description of the video — include camera movement, lighting,
369
+ subject action, and mood for best results.
370
+ duration_seconds: Video duration in seconds (typically 5–8).
371
+ aspect_ratio: '16:9' for landscape or '9:16' for portrait/mobile.
372
+ negative_prompt: What to exclude from the video.
373
+ anchor_frame_path: Optional path to an image to use as the first frame.
374
+ enhance_prompt: Whether to allow Veo to rewrite/enhance your prompt.
375
+ output_dir: Override the output directory (default: ~/google_flow_outputs).
376
+
377
+ Returns:
378
+ Paths or URIs to the generated video(s).
379
+ """
380
+ client = get_client()
381
+ out = Path(output_dir) if output_dir else _output_dir()
382
+ out.mkdir(parents=True, exist_ok=True)
383
+
384
+ config = types.GenerateVideosConfig(
385
+ duration_seconds=duration_seconds,
386
+ aspect_ratio=aspect_ratio,
387
+ negative_prompt=negative_prompt,
388
+ enhance_prompt=enhance_prompt,
389
+ )
390
+
391
+ kwargs: dict = dict(model=MODEL_VIDEO, prompt=prompt, config=config)
392
+
393
+ if anchor_frame_path:
394
+ kwargs["image"] = types.Image(
395
+ image_bytes=_load_image_bytes(anchor_frame_path),
396
+ mime_type=_mime(anchor_frame_path),
397
+ )
398
+
399
+ operation = await client.aio.models.generate_videos(**kwargs)
400
+ operation = await _poll_operation(client, operation)
401
+
402
+ ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
403
+ saved, uris = _collect_video_results(operation, out, ts)
404
+ return _format_video_result(saved, uris)
405
+
406
+
407
+ # ── Tool 5: Extend Video ──────────────────────────────────────────────────────
408
+
409
+
410
+ @mcp.tool()
411
+ async def flow_extend_video(
412
+ video_path: str,
413
+ prompt: str,
414
+ duration_seconds: int = 8,
415
+ aspect_ratio: str = "16:9",
416
+ output_dir: Optional[str] = None,
417
+ ) -> str:
418
+ """
419
+ Extend an existing Veo 3.1 video clip with additional generated content.
420
+
421
+ Requires a paid Google AI API plan. The extended video is SynthID-watermarked.
422
+
423
+ Args:
424
+ video_path: Absolute path to the source MP4 video file to extend.
425
+ prompt: Description of how to continue the video after the source clip ends.
426
+ duration_seconds: Duration of the generated extension in seconds (typically 5–8).
427
+ aspect_ratio: Must match the source video — '16:9' or '9:16'.
428
+ output_dir: Override the output directory (default: ~/google_flow_outputs).
429
+
430
+ Returns:
431
+ Path or URI to the extended video.
432
+ """
433
+ client = get_client()
434
+ out = Path(output_dir) if output_dir else _output_dir()
435
+ out.mkdir(parents=True, exist_ok=True)
436
+
437
+ with open(video_path, "rb") as f:
438
+ video_bytes = f.read()
439
+
440
+ operation = await client.aio.models.generate_videos(
441
+ model=MODEL_VIDEO,
442
+ prompt=prompt,
443
+ video=types.Video(video_bytes=video_bytes, mime_type="video/mp4"),
444
+ config=types.GenerateVideosConfig(
445
+ duration_seconds=duration_seconds,
446
+ aspect_ratio=aspect_ratio,
447
+ ),
448
+ )
449
+ operation = await _poll_operation(client, operation)
450
+
451
+ ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
452
+ saved, uris = _collect_video_results(operation, out, ts)
453
+ return _format_video_result(saved, uris)
454
+
455
+
456
+ # ── Tool 6: Image → Video Pipeline ───────────────────────────────────────────
457
+
458
+
459
+ @mcp.tool()
460
+ async def flow_image_to_video(
461
+ video_prompt: str,
462
+ image_prompt: Optional[str] = None,
463
+ image_path: Optional[str] = None,
464
+ image_model: str = "pro",
465
+ duration_seconds: int = 8,
466
+ aspect_ratio: str = "16:9",
467
+ negative_prompt: Optional[str] = None,
468
+ output_dir: Optional[str] = None,
469
+ ) -> str:
470
+ """
471
+ Full pipeline: generate or use an image with Nano Banana Pro, then animate it with Veo 3.1.
472
+
473
+ Provide either image_prompt (to generate a new anchor image) or image_path (to use an
474
+ existing image). The anchor image becomes the first frame of the video.
475
+
476
+ Requires a paid Google AI API plan for the video step.
477
+
478
+ Args:
479
+ video_prompt: Cinematic description of the video animation — how the scene moves.
480
+ image_prompt: Text prompt to generate the anchor image (if image_path is not given).
481
+ image_path: Path to an existing image to use as the anchor frame.
482
+ image_model: 'pro' (Nano Banana Pro, recommended) or 'flash' (Nano Banana 2).
483
+ duration_seconds: Video duration in seconds (typically 5–8).
484
+ aspect_ratio: '16:9' for landscape or '9:16' for portrait.
485
+ negative_prompt: What to exclude from both the image and video.
486
+ output_dir: Override the output directory (default: ~/google_flow_outputs).
487
+
488
+ Returns:
489
+ Paths to the generated anchor image and the final video.
490
+ """
491
+ if not image_prompt and not image_path:
492
+ raise ValueError("Provide either image_prompt or image_path.")
493
+
494
+ client = get_client()
495
+ out = Path(output_dir) if output_dir else _output_dir()
496
+ out.mkdir(parents=True, exist_ok=True)
497
+ results: list[str] = []
498
+
499
+ # Step 1 — Generate or load the anchor image
500
+ if image_path:
501
+ anchor_path = image_path
502
+ else:
503
+ model_id = _resolve_image_model(image_model)
504
+ img_response = await client.aio.models.generate_images(
505
+ model=model_id,
506
+ prompt=image_prompt,
507
+ config=types.GenerateImagesConfig(
508
+ number_of_images=1,
509
+ aspect_ratio=aspect_ratio,
510
+ negative_prompt=negative_prompt,
511
+ ),
512
+ )
513
+ ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
514
+ anchor = out / f"flow_anchor_{ts}.png"
515
+ _save_image(img_response.generated_images[0].image, anchor)
516
+ anchor_path = str(anchor)
517
+ results.append(f"Anchor image:\n{anchor_path}")
518
+
519
+ # Step 2 — Animate with Veo 3.1
520
+ operation = await client.aio.models.generate_videos(
521
+ model=MODEL_VIDEO,
522
+ prompt=video_prompt,
523
+ image=types.Image(
524
+ image_bytes=_load_image_bytes(anchor_path),
525
+ mime_type=_mime(anchor_path),
526
+ ),
527
+ config=types.GenerateVideosConfig(
528
+ duration_seconds=duration_seconds,
529
+ aspect_ratio=aspect_ratio,
530
+ negative_prompt=negative_prompt,
531
+ ),
532
+ )
533
+ operation = await _poll_operation(client, operation)
534
+
535
+ ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
536
+ saved, uris = _collect_video_results(operation, out, ts)
537
+ results.append(_format_video_result(saved, uris))
538
+
539
+ return "\n\n".join(results)
540
+
541
+
542
+ # ── Entry Point ───────────────────────────────────────────────────────────────
543
+
544
+
545
+ def main() -> None:
546
+ """Entry point for uvx / python -m google_flow_mcp."""
547
+ mcp.run(transport="stdio")
548
+
549
+
550
+ if __name__ == "__main__":
551
+ main()
File without changes
@@ -0,0 +1,252 @@
1
+ """Unit tests for google_flow_mcp.server — all API calls are mocked."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import os
6
+ import tempfile
7
+ from pathlib import Path
8
+ from unittest.mock import AsyncMock, MagicMock
9
+
10
+ import pytest
11
+
12
+ # Set env vars before importing the server module
13
+ os.environ.setdefault("GOOGLE_API_KEY", "test-api-key-000")
14
+ os.environ["FLOW_OUTPUT_DIR"] = tempfile.mkdtemp()
15
+
16
+ import google_flow_mcp.server as server_module # noqa: E402
17
+ from google_flow_mcp.server import ( # noqa: E402
18
+ MODEL_IMAGE_FLASH,
19
+ MODEL_IMAGE_PRO,
20
+ MODEL_VIDEO,
21
+ flow_edit_image,
22
+ flow_extend_video,
23
+ flow_generate_image,
24
+ flow_generate_image_with_references,
25
+ flow_generate_video,
26
+ flow_image_to_video,
27
+ )
28
+
29
+
30
+ # ── Fixtures ──────────────────────────────────────────────────────────────────
31
+
32
+
33
+ @pytest.fixture(autouse=True)
34
+ def reset_singleton():
35
+ """Reset the lazy client singleton between tests."""
36
+ server_module._client = None
37
+ yield
38
+ server_module._client = None
39
+
40
+
41
+ @pytest.fixture
42
+ def mock_client(mocker) -> MagicMock:
43
+ """Patch _get_client to return a fully mocked async client."""
44
+ client = MagicMock()
45
+ client.aio = MagicMock()
46
+ client.aio.models = MagicMock()
47
+ client.aio.models.generate_images = AsyncMock()
48
+ client.aio.models.edit_image = AsyncMock()
49
+ client.aio.models.generate_videos = AsyncMock()
50
+ client.aio.operations = MagicMock()
51
+ client.aio.operations.get = AsyncMock()
52
+ mocker.patch.object(server_module, "_get_client", return_value=client)
53
+ return client
54
+
55
+
56
+ def _fake_gen_image(tmp_path: Path, name: str = "out.png") -> MagicMock:
57
+ """Build a mock GeneratedImage whose .image.save() actually creates a file."""
58
+ path_ref: list[Path] = []
59
+ target = tmp_path / name
60
+
61
+ gen = MagicMock()
62
+ gen.image.image_bytes = None
63
+
64
+ def save_side(path_str: str) -> None:
65
+ p = Path(path_str)
66
+ path_ref.append(p)
67
+ p.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 32)
68
+
69
+ gen.image.save.side_effect = save_side
70
+ gen.image._save_target = target
71
+ return gen
72
+
73
+
74
+ def _fake_video_operation(uri: str = "gs://bucket/out.mp4", done: bool = True) -> MagicMock:
75
+ """Build a mock Veo operation that is immediately done."""
76
+ op = MagicMock()
77
+ op.done = done
78
+ op.error = MagicMock()
79
+ op.error.message = None
80
+ vid = MagicMock()
81
+ vid.video.video_bytes = None
82
+ vid.video.uri = uri
83
+ op.response.generated_videos = [vid]
84
+ return op
85
+
86
+
87
+ def _write_tmp_png(tmp_path: Path, name: str = "input.png") -> str:
88
+ """Write a minimal PNG file and return its path as string."""
89
+ p = tmp_path / name
90
+ p.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 32)
91
+ return str(p)
92
+
93
+
94
+ def _write_tmp_mp4(tmp_path: Path, name: str = "input.mp4") -> str:
95
+ p = tmp_path / name
96
+ p.write_bytes(b"\x00\x00\x00\x18ftypmp42" + b"\x00" * 32)
97
+ return str(p)
98
+
99
+
100
+ # ── Tests ─────────────────────────────────────────────────────────────────────
101
+
102
+
103
+ @pytest.mark.asyncio
104
+ async def test_01_generate_image_default_model(mock_client, tmp_path):
105
+ """flow_generate_image uses MODEL_IMAGE_PRO by default."""
106
+ gen = _fake_gen_image(tmp_path)
107
+ mock_client.aio.models.generate_images.return_value.generated_images = [gen]
108
+
109
+ result = await flow_generate_image(prompt="a red apple", output_dir=str(tmp_path))
110
+
111
+ call_kwargs = mock_client.aio.models.generate_images.call_args.kwargs
112
+ assert call_kwargs["model"] == MODEL_IMAGE_PRO
113
+ assert call_kwargs["prompt"] == "a red apple"
114
+ assert "flow_image_" in result
115
+
116
+
117
+ @pytest.mark.asyncio
118
+ async def test_02_generate_image_flash_model(mock_client, tmp_path):
119
+ """Passing model='flash' selects MODEL_IMAGE_FLASH."""
120
+ gen = _fake_gen_image(tmp_path)
121
+ mock_client.aio.models.generate_images.return_value.generated_images = [gen]
122
+
123
+ await flow_generate_image(prompt="quick sketch", model="flash", output_dir=str(tmp_path))
124
+
125
+ call_kwargs = mock_client.aio.models.generate_images.call_args.kwargs
126
+ assert call_kwargs["model"] == MODEL_IMAGE_FLASH
127
+
128
+
129
+ @pytest.mark.asyncio
130
+ async def test_03_generate_image_count_clamped_to_4(mock_client, tmp_path):
131
+ """number_of_images is clamped to max 4 before being sent to the API."""
132
+ gens = [_fake_gen_image(tmp_path, f"img{i}.png") for i in range(4)]
133
+ mock_client.aio.models.generate_images.return_value.generated_images = gens
134
+
135
+ result = await flow_generate_image(
136
+ prompt="many apples", number_of_images=99, output_dir=str(tmp_path)
137
+ )
138
+
139
+ config_arg = mock_client.aio.models.generate_images.call_args.kwargs["config"]
140
+ assert config_arg.number_of_images == 4
141
+ assert "4 image(s)" in result
142
+
143
+
144
+ @pytest.mark.asyncio
145
+ async def test_04_generate_image_multiple_files_saved(mock_client, tmp_path):
146
+ """Result lists all saved paths when multiple images are returned."""
147
+ gens = [_fake_gen_image(tmp_path, f"r{i}.png") for i in range(3)]
148
+ mock_client.aio.models.generate_images.return_value.generated_images = gens
149
+
150
+ result = await flow_generate_image(
151
+ prompt="trio", number_of_images=3, output_dir=str(tmp_path)
152
+ )
153
+
154
+ assert "3 image(s)" in result
155
+ assert result.count("flow_image_") == 3
156
+
157
+
158
+ @pytest.mark.asyncio
159
+ async def test_05_edit_image_passes_reference(mock_client, tmp_path):
160
+ """flow_edit_image passes a RawReferenceImage to generate_images."""
161
+ src = _write_tmp_png(tmp_path, "src.png")
162
+ gen = _fake_gen_image(tmp_path, "edited.png")
163
+ mock_client.aio.models.edit_image.return_value.generated_images = [gen]
164
+
165
+ result = await flow_edit_image(
166
+ image_path=src, prompt="make it blue", output_dir=str(tmp_path)
167
+ )
168
+
169
+ call_kwargs = mock_client.aio.models.edit_image.call_args.kwargs
170
+ ref_images = call_kwargs["reference_images"]
171
+ assert len(ref_images) == 2 # RawReferenceImage + MaskReferenceImage
172
+ assert "flow_edit_" in result
173
+
174
+
175
+ @pytest.mark.asyncio
176
+ async def test_06_generate_with_references_passes_all_refs(mock_client, tmp_path):
177
+ """flow_generate_image_with_references passes one RawReferenceImage per path."""
178
+ paths = [_write_tmp_png(tmp_path, f"ref{i}.png") for i in range(3)]
179
+ gen = _fake_gen_image(tmp_path, "out.png")
180
+ mock_client.aio.models.generate_images.return_value.generated_images = [gen]
181
+
182
+ result = await flow_generate_image_with_references(
183
+ prompt="inspired by these", reference_image_paths=paths, output_dir=str(tmp_path)
184
+ )
185
+
186
+ call_kwargs = mock_client.aio.models.generate_images.call_args.kwargs
187
+ assert len(call_kwargs["reference_images"]) == 3
188
+ assert "3 reference(s)" in result
189
+
190
+
191
+ @pytest.mark.asyncio
192
+ async def test_07_generate_with_references_too_many_raises(mock_client, tmp_path):
193
+ """More than 14 reference images raises ValueError."""
194
+ paths = [_write_tmp_png(tmp_path, f"r{i}.png") for i in range(15)]
195
+
196
+ with pytest.raises(ValueError, match="14"):
197
+ await flow_generate_image_with_references(
198
+ prompt="too many", reference_image_paths=paths, output_dir=str(tmp_path)
199
+ )
200
+
201
+
202
+ @pytest.mark.asyncio
203
+ async def test_08_generate_video_success(mock_client, tmp_path):
204
+ """flow_generate_video calls generate_videos with MODEL_VIDEO and returns URI."""
205
+ op = _fake_video_operation(uri="gs://bucket/vid.mp4")
206
+ mock_client.aio.models.generate_videos.return_value = op
207
+
208
+ result = await flow_generate_video(prompt="drone over mountains", output_dir=str(tmp_path))
209
+
210
+ call_kwargs = mock_client.aio.models.generate_videos.call_args.kwargs
211
+ assert call_kwargs["model"] == MODEL_VIDEO
212
+ assert "SynthID" in result
213
+
214
+
215
+ @pytest.mark.asyncio
216
+ async def test_09_generate_video_polls_until_done(mock_client, tmp_path):
217
+ """flow_generate_video polls the operation when initially not done."""
218
+ op_pending = _fake_video_operation(done=False)
219
+ op_done = _fake_video_operation(done=True, uri="gs://bucket/final.mp4")
220
+ mock_client.aio.models.generate_videos.return_value = op_pending
221
+ mock_client.aio.operations.get.return_value = op_done
222
+
223
+ # Patch sleep to avoid actual delay
224
+ import asyncio
225
+ from unittest.mock import patch
226
+
227
+ with patch.object(asyncio, "sleep", new_callable=AsyncMock):
228
+ result = await flow_generate_video(prompt="slow sunrise", output_dir=str(tmp_path))
229
+
230
+ mock_client.aio.operations.get.assert_called_once_with(op_pending)
231
+ assert "SynthID" in result
232
+
233
+
234
+ @pytest.mark.asyncio
235
+ async def test_10_image_to_video_generates_anchor_then_video(mock_client, tmp_path):
236
+ """flow_image_to_video calls generate_images first, then generate_videos."""
237
+ gen = _fake_gen_image(tmp_path, "anchor.png")
238
+ mock_client.aio.models.generate_images.return_value.generated_images = [gen]
239
+
240
+ op = _fake_video_operation(uri="gs://bucket/final.mp4")
241
+ mock_client.aio.models.generate_videos.return_value = op
242
+
243
+ result = await flow_image_to_video(
244
+ video_prompt="the scene comes alive",
245
+ image_prompt="a tranquil mountain lake at dawn",
246
+ output_dir=str(tmp_path),
247
+ )
248
+
249
+ mock_client.aio.models.generate_images.assert_called_once()
250
+ mock_client.aio.models.generate_videos.assert_called_once()
251
+ assert "Anchor image" in result
252
+ assert "SynthID" in result