475 lines
20 KiB
Python
475 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Generates and edits videos using the Gemini Omni Flash model via the google-genai Interactions API.
|
|
Can automatically upload local media references using the Files API.
|
|
Supports parallel execution of multiple generations using Python standard library.
|
|
Uses the official google-genai SDK.
|
|
"""
|
|
|
|
import argparse
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
import urllib.error
|
|
import urllib.parse
|
|
import uuid
|
|
from google import genai
|
|
|
|
# Load local upload helper logic inline to prevent dependency issues
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from upload_file import upload_file, wait_for_active
|
|
from pathlib import Path
|
|
|
|
|
|
def safe_user_path(path_value, base_dir="."):
|
|
"""Resolve a CLI path under the current workspace."""
|
|
if base_dir != ".":
|
|
raise ValueError("Custom base directories are not supported for CLI paths")
|
|
base_path = Path.cwd().resolve()
|
|
resolved_path = Path(path_value).expanduser().resolve()
|
|
try:
|
|
resolved_path.relative_to(base_path)
|
|
except ValueError as exc:
|
|
raise ValueError(f"Path escapes allowed directory: {path_value}") from exc
|
|
return resolved_path
|
|
|
|
def get_api_key(args):
|
|
"""Retrieves API key from command args or environment."""
|
|
if args.api_key:
|
|
return args.api_key
|
|
return os.environ.get("GEMINI_API_KEY")
|
|
|
|
FILE_ID_RE = re.compile(r'^[A-Za-z0-9_-]+$')
|
|
|
|
|
|
def extract_file_id(uri):
|
|
"""Returns a Gemini File API id from a local reference or trusted API URL."""
|
|
if not uri:
|
|
return None
|
|
if uri.startswith("files/"):
|
|
file_id = uri.removeprefix("files/")
|
|
return file_id if FILE_ID_RE.fullmatch(file_id) else None
|
|
|
|
parsed = urllib.parse.urlparse(uri)
|
|
if parsed.scheme != "https" or parsed.netloc != "generativelanguage.googleapis.com":
|
|
return None
|
|
path_match = re.fullmatch(r'/files/([A-Za-z0-9_-]+)', parsed.path)
|
|
return path_match.group(1) if path_match else None
|
|
|
|
|
|
def is_file_uri(uri):
|
|
"""Returns True if the string is a standard Gemini File URI."""
|
|
return extract_file_id(uri) is not None
|
|
|
|
def normalize_file_uri(uri):
|
|
"""Normalizes any File API URI/reference to the standard https://generativelanguage.googleapis.com/files/{id} format."""
|
|
file_id = extract_file_id(uri)
|
|
if file_id:
|
|
return f"https://generativelanguage.googleapis.com/files/{file_id}"
|
|
return uri
|
|
|
|
|
|
def media_download_url(file_uri):
|
|
"""Build a media URL only for validated Gemini File API references."""
|
|
file_id = extract_file_id(file_uri)
|
|
if not file_id:
|
|
raise ValueError("Generated video URI must be a Gemini File API reference.")
|
|
return f"https://generativelanguage.googleapis.com/files/{file_id}?alt=media"
|
|
|
|
def slugify(text):
|
|
"""Converts a text prompt into a safe, descriptive filename slug."""
|
|
text = text.lower()
|
|
text = re.sub(r'[^a-z0-9]+', '_', text)
|
|
return text.strip('_')[:50]
|
|
|
|
def parse_and_validate_duration(value):
|
|
"""Parses and formats a duration integer between 3 and 10 with optional 's' suffix."""
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, (int, float)):
|
|
val = float(value)
|
|
else:
|
|
clean_value = str(value).strip().lower()
|
|
if clean_value in ('none', ''):
|
|
return None
|
|
if clean_value.endswith('s'):
|
|
clean_value = clean_value[:-1]
|
|
try:
|
|
val = float(clean_value)
|
|
except ValueError:
|
|
raise ValueError(f"Invalid duration value: '{value}'. Must be an integer (e.g., 5, 10).")
|
|
|
|
if not val.is_integer():
|
|
raise ValueError(f"Duration must be an integer, not a float (e.g., got {value}).")
|
|
|
|
val_int = int(val)
|
|
if val_int < 3 or val_int > 10:
|
|
raise ValueError(f"Duration must be between 3 (inclusive) and 10 (inclusive) seconds. Got {val_int}.")
|
|
|
|
return f"{val_int}s"
|
|
|
|
def argparse_duration_type(value):
|
|
"""argparse type converter for validating duration."""
|
|
if value is None or str(value).strip().lower() in ('none', ''):
|
|
return None
|
|
try:
|
|
return parse_and_validate_duration(value)
|
|
except ValueError as e:
|
|
raise argparse.ArgumentTypeError(str(e))
|
|
|
|
def resolve_or_upload_asset(asset_path, mime_type, api_key, strip_audio=False):
|
|
"""
|
|
If asset_path is a File API URI, returns it directly (normalized).
|
|
If it is a local file path, uploads it and returns its File API URI (normalized).
|
|
"""
|
|
if not asset_path:
|
|
return None, None
|
|
|
|
if is_file_uri(asset_path):
|
|
normalized = normalize_file_uri(asset_path)
|
|
print(f"Using existing File URI: {normalized}")
|
|
if strip_audio:
|
|
print("Warning: --strip-audio was specified but the video input is an existing remote File URI. "
|
|
"Audio cannot be stripped from remote files automatically.")
|
|
return normalized, mime_type
|
|
|
|
if os.path.exists(asset_path):
|
|
upload_path = asset_path
|
|
temp_stripped_path = None
|
|
|
|
if strip_audio:
|
|
print(f"Detected local asset path '{asset_path}'. Stripping audio before upload...")
|
|
|
|
# Check if ffmpeg is available
|
|
import subprocess
|
|
try:
|
|
subprocess.run(["ffmpeg", "-version"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
|
|
except (subprocess.SubprocessError, FileNotFoundError):
|
|
raise RuntimeError(
|
|
"Error: ffmpeg is not installed or not found in system PATH. "
|
|
"ffmpeg is required to strip audio from local videos."
|
|
)
|
|
|
|
try:
|
|
os.makedirs("media", exist_ok=True)
|
|
base_name = os.path.basename(asset_path)
|
|
name, ext = os.path.splitext(base_name)
|
|
temp_stripped_path = os.path.join("media", f"temp_stripped_{name}_{uuid.uuid4().hex}{ext}")
|
|
|
|
# Fast stream-copy audio stripping
|
|
cmd = ["ffmpeg", "-y", "-i", asset_path, "-c:v", "copy", "-an", temp_stripped_path]
|
|
subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
|
|
|
|
print(f"Successfully stripped audio. Temporary video file created at: {temp_stripped_path}")
|
|
upload_path = temp_stripped_path
|
|
except Exception as e:
|
|
print(f"Error stripping audio using ffmpeg: {e}", file=sys.stderr)
|
|
print("Falling back to uploading the original video with audio.", file=sys.stderr)
|
|
|
|
print(f"Uploading asset '{upload_path}'...")
|
|
file_meta = upload_file(upload_path, api_key=api_key)
|
|
file_name = file_meta.get("name")
|
|
# Wait for file to become active
|
|
file_meta = wait_for_active(file_name, api_key)
|
|
normalized = normalize_file_uri(file_meta.get("uri"))
|
|
|
|
# Clean up temporary stripped file if we created one
|
|
if temp_stripped_path and os.path.exists(temp_stripped_path):
|
|
try:
|
|
safe_user_path(temp_stripped_path).unlink()
|
|
print(f"Cleaned up temporary video file: {temp_stripped_path}")
|
|
except Exception as e:
|
|
print(f"Warning: Failed to remove temporary file {temp_stripped_path}: {e}", file=sys.stderr)
|
|
|
|
# Handle both mimeType and mime_type key formats returned from upload_file
|
|
returned_mime = file_meta.get("mimeType") or file_meta.get("mime_type")
|
|
return normalized, returned_mime
|
|
else:
|
|
raise FileNotFoundError(f"Asset path '{asset_path}' is neither a valid File API URI nor a local file path.")
|
|
|
|
def download_video_file(file_uri, output_path, api_key):
|
|
"""Downloads generated video file from URI using alt=media standard in a memory-safe, chunked manner."""
|
|
download_url = media_download_url(file_uri)
|
|
|
|
print(f"Downloading video from {file_uri} to {output_path} in chunked mode...")
|
|
req = urllib.request.Request(download_url)
|
|
req.add_header("x-goog-api-key", api_key)
|
|
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=480) as resp:
|
|
parent_dir = os.path.dirname(output_path)
|
|
if parent_dir:
|
|
os.makedirs(parent_dir, exist_ok=True)
|
|
|
|
with safe_user_path(output_path).open("wb") as f:
|
|
while True:
|
|
chunk = resp.read(8192)
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
print(f"Video successfully saved to: {output_path}")
|
|
except urllib.error.HTTPError as e:
|
|
raise RuntimeError(f"Error downloading video file: {e.code} - {e.read().decode()}")
|
|
|
|
def generate_video(prompt, api_key, model="gemini-omni-flash-preview", aspect_ratio="16:9", duration=None, image_path=None, video_path=None, output_path="output.mp4", strip_audio=False, previous_interaction_id=None):
|
|
"""Creates an interaction with the video model and downloads the resulting video using the official google-genai SDK."""
|
|
duration = parse_and_validate_duration(duration)
|
|
input_parts = []
|
|
|
|
# 1. Resolve and add image inputs (reference/start/end frames)
|
|
if image_path:
|
|
if isinstance(image_path, list):
|
|
for path in image_path:
|
|
img_uri, img_mime = resolve_or_upload_asset(path, "image/png", api_key)
|
|
input_parts.append({
|
|
"type": "image",
|
|
"uri": img_uri,
|
|
"mime_type": img_mime
|
|
})
|
|
else:
|
|
img_uri, img_mime = resolve_or_upload_asset(image_path, "image/png", api_key)
|
|
input_parts.append({
|
|
"type": "image",
|
|
"uri": img_uri,
|
|
"mime_type": img_mime
|
|
})
|
|
|
|
# 2. Resolve and add video inputs (for edits or extensions)
|
|
if video_path:
|
|
if isinstance(video_path, list):
|
|
for path in video_path:
|
|
vid_uri, vid_mime = resolve_or_upload_asset(path, "video/mp4", api_key, strip_audio=strip_audio)
|
|
input_parts.append({
|
|
"type": "video",
|
|
"uri": vid_uri,
|
|
"mime_type": vid_mime
|
|
})
|
|
else:
|
|
vid_uri, vid_mime = resolve_or_upload_asset(video_path, "video/mp4", api_key, strip_audio=strip_audio)
|
|
input_parts.append({
|
|
"type": "video",
|
|
"uri": vid_uri,
|
|
"mime_type": vid_mime
|
|
})
|
|
|
|
# 3. Add text prompt
|
|
input_parts.append({
|
|
"type": "text",
|
|
"text": prompt
|
|
})
|
|
|
|
# Construct the config
|
|
video_config = {
|
|
"type": "video",
|
|
"aspect_ratio": aspect_ratio,
|
|
"delivery": "uri"
|
|
}
|
|
if duration:
|
|
video_config["duration"] = duration
|
|
|
|
print(f"\nSending generation request using official google-genai SDK and model '{model}'...")
|
|
print(f"Prompt: '{prompt}' | Aspect Ratio: {aspect_ratio} | Duration: {duration}")
|
|
|
|
# Initialize the client and call interactions.create
|
|
client = genai.Client(api_key=api_key)
|
|
try:
|
|
interaction = client.interactions.create(
|
|
model=model,
|
|
input=input_parts,
|
|
response_format=video_config,
|
|
previous_interaction_id=previous_interaction_id
|
|
)
|
|
except Exception as e:
|
|
raise RuntimeError(f"Error generating video via SDK: {e}")
|
|
|
|
print(f"Generation complete for '{prompt}'! Processing response...")
|
|
|
|
interaction_id = interaction.id
|
|
if interaction_id:
|
|
print(f"Interaction ID: {interaction_id}")
|
|
|
|
output_video = interaction.output_video
|
|
if not output_video or not output_video.uri:
|
|
err_msg = f"No video content found in response for '{prompt}'."
|
|
if video_path:
|
|
err_msg += (
|
|
"\nWARNING: IMPORTANT REGIONAL RESTRICTION: Uploading videos to use for video edits is "
|
|
"not available in the EEA, Switzerland, United Kingdom, and some US states."
|
|
)
|
|
raise RuntimeError(f"{err_msg}\nResponse output_video field: {output_video}")
|
|
|
|
video_uri = output_video.uri
|
|
print(f"Generated video URI for '{prompt}': {video_uri}")
|
|
|
|
# Download the final video
|
|
download_video_file(video_uri, output_path, api_key)
|
|
|
|
def run_job(job, api_key):
|
|
"""Runs a single generation job inside a thread pool, catching exceptions."""
|
|
prompt = job.get("prompt")
|
|
if not prompt:
|
|
print("Warning: Skipping job with empty prompt.", file=sys.stderr)
|
|
return {"job": job, "status": "SKIPPED", "error": "Empty prompt"}
|
|
|
|
aspect_ratio = job.get("aspect_ratio", "16:9")
|
|
duration = job.get("duration")
|
|
image_path = job.get("image")
|
|
video_path = job.get("video")
|
|
output_path = job.get("output")
|
|
model = job.get("model", "gemini-omni-flash-preview")
|
|
strip_audio = job.get("strip_audio", False)
|
|
previous_interaction_id = job.get("previous_interaction_id")
|
|
|
|
if not output_path:
|
|
output_path = f"media/output_{slugify(prompt)}.mp4"
|
|
|
|
print(f"[Parallel] Dispatching: '{prompt}' (Output: {output_path})")
|
|
|
|
try:
|
|
generate_video(
|
|
prompt=prompt,
|
|
api_key=api_key,
|
|
model=model,
|
|
aspect_ratio=aspect_ratio,
|
|
duration=duration,
|
|
image_path=image_path,
|
|
video_path=video_path,
|
|
output_path=output_path,
|
|
strip_audio=strip_audio,
|
|
previous_interaction_id=previous_interaction_id
|
|
)
|
|
return {"job": job, "status": "SUCCESS", "output_path": output_path}
|
|
except Exception as e:
|
|
print(f"[Parallel] Failed: '{prompt}' - Error: {e}", file=sys.stderr)
|
|
return {"job": job, "status": "FAILED", "error": str(e)}
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Generate and edit videos using Gemini Omni Flash model via google-genai SDK (supports parallel batch execution).")
|
|
parser.add_argument("prompt", nargs="?", help="Text prompt / instruction for a single video generation")
|
|
parser.add_argument("--image", action="append", help="Optional local image path or File API URI for referencing / image-to-video (can be specified multiple times)")
|
|
parser.add_argument("--video", action="append", help="Optional local video path or File API URI for editing / extending (can be specified multiple times)")
|
|
parser.add_argument("--aspect-ratio", default="16:9", choices=["16:9", "9:16"], help="Aspect ratio (default: 16:9)")
|
|
parser.add_argument("--duration", type=argparse_duration_type, default=None, help="Video duration as an integer between 3 and 10 seconds (e.g., 5, 10). Default: None (API/Model decides, typically 10s or matches source)")
|
|
parser.add_argument("--model", default="gemini-omni-flash-preview", help="Gemini Omni Flash video model ID (default: gemini-omni-flash-preview)")
|
|
parser.add_argument("--output", help="Local output file path for single generation (default: media/output.mp4)")
|
|
parser.add_argument("--strip-audio", "-a", action="store_true", help="Completely strip/disable audio stream from the input video(s) before uploading so Gemini Omni Flash can regenerate new audio from scratch")
|
|
parser.add_argument("--previous-interaction-id", help="Optional Interaction ID of a previous generation for turn-by-turn editing")
|
|
parser.add_argument("--api-key", help="Gemini API Key (overrides env)")
|
|
|
|
# Parallel batch configuration options
|
|
parser.add_argument("--batch", help="Path to a JSON file containing an array of generation jobs")
|
|
parser.add_argument("--prompts-file", help="Path to a text file containing one prompt per line to run in parallel")
|
|
parser.add_argument("--concurrency", type=int, default=3, help="Maximum number of concurrent executions (default: 3)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
api_key = get_api_key(args)
|
|
if not api_key:
|
|
print("Error: API key is not set. Use --api-key or set GEMINI_API_KEY environment variable.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# 1. Handle Batch JSON execution
|
|
if args.batch:
|
|
if not os.path.exists(args.batch):
|
|
print(f"Error: Batch JSON file '{args.batch}' not found.", file=sys.stderr)
|
|
sys.exit(1)
|
|
try:
|
|
with safe_user_path(args.batch).open("r", encoding="utf-8") as f:
|
|
jobs = json.load(f)
|
|
if not isinstance(jobs, list):
|
|
print("Error: Batch JSON file must contain a list/array of job objects.", file=sys.stderr)
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(f"Error parsing Batch JSON: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
print(f"Loaded {len(jobs)} jobs from batch JSON. Running with concurrency={args.concurrency}...")
|
|
|
|
# 2. Handle Prompts File execution
|
|
elif args.prompts_file:
|
|
if not os.path.exists(args.prompts_file):
|
|
print(f"Error: Prompts file '{args.prompts_file}' not found.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
jobs = []
|
|
with safe_user_path(args.prompts_file).open("r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line and not line.startswith("#"):
|
|
jobs.append({
|
|
"prompt": line,
|
|
"aspect_ratio": args.aspect_ratio,
|
|
"duration": args.duration,
|
|
"image": args.image,
|
|
"video": args.video,
|
|
"model": args.model,
|
|
"strip_audio": args.strip_audio,
|
|
"previous_interaction_id": args.previous_interaction_id
|
|
})
|
|
print(f"Loaded {len(jobs)} prompts from text file. Running with concurrency={args.concurrency}...")
|
|
|
|
# 3. Handle standard single prompt execution
|
|
else:
|
|
if not args.prompt:
|
|
parser.print_help()
|
|
sys.exit(1)
|
|
|
|
output_path = args.output if args.output else "media/output.mp4"
|
|
try:
|
|
generate_video(
|
|
prompt=args.prompt,
|
|
api_key=api_key,
|
|
model=args.model,
|
|
aspect_ratio=args.aspect_ratio,
|
|
duration=args.duration,
|
|
image_path=args.image,
|
|
video_path=args.video,
|
|
output_path=output_path,
|
|
strip_audio=args.strip_audio,
|
|
previous_interaction_id=args.previous_interaction_id
|
|
)
|
|
sys.exit(0)
|
|
except Exception as e:
|
|
print(f"Error: Generation failed: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# Parallel Execution Loop
|
|
if not jobs:
|
|
print("Warning: No valid jobs found to execute.")
|
|
sys.exit(0)
|
|
|
|
results = []
|
|
with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
|
|
futures = {executor.submit(run_job, job, api_key): job for job in jobs}
|
|
for future in as_completed(futures):
|
|
results.append(future.result())
|
|
|
|
# Print Batch Results Summary
|
|
print("\n" + "="*50)
|
|
print("BATCH PARALLEL EXECUTION SUMMARY")
|
|
print("="*50)
|
|
success_count = sum(1 for r in results if r["status"] == "SUCCESS")
|
|
failed_count = sum(1 for r in results if r["status"] == "FAILED")
|
|
skipped_count = sum(1 for r in results if r["status"] == "SKIPPED")
|
|
|
|
print(f"Total: {len(results)} | Success: {success_count} | Failed: {failed_count} | Skipped: {skipped_count}\n")
|
|
for r in results:
|
|
status_str = r["status"]
|
|
prompt = r["job"].get("prompt")
|
|
if r["status"] == "SUCCESS":
|
|
print(f" [{status_str}] '{prompt}' -> {r['output_path']}")
|
|
else:
|
|
print(f" [{status_str}] '{prompt}' -> Error: {r.get('error')}")
|
|
print("="*50)
|
|
|
|
if failed_count > 0:
|
|
sys.exit(1)
|
|
sys.exit(0)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|