103 lines
3.0 KiB
Python
103 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate images via the Grok Build CLI /imagine command.
|
|
|
|
Requires an authenticated grok CLI (SuperGrok/X Premium+ subscription):
|
|
grok login
|
|
"""
|
|
|
|
import argparse
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
|
|
GROK = shutil.which("grok") or str(Path.home() / ".grok" / "bin" / "grok")
|
|
SESSION_ROOT = Path.home() / ".grok" / "sessions"
|
|
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp"}
|
|
|
|
|
|
def die(msg: str) -> None:
|
|
print(f"error: {msg}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
def newest_image(since_ts: float) -> Optional[Path]:
|
|
best = None
|
|
if not SESSION_ROOT.is_dir():
|
|
return None
|
|
for p in SESSION_ROOT.rglob("*"):
|
|
if not p.is_file() or p.suffix.lower() not in IMAGE_EXTS:
|
|
continue
|
|
try:
|
|
mtime = p.stat().st_mtime
|
|
except OSError:
|
|
continue
|
|
if mtime >= since_ts and (best is None or mtime > best.stat().st_mtime):
|
|
best = p
|
|
return best
|
|
|
|
|
|
def path_from_output(text: str) -> Optional[Path]:
|
|
m = re.search(r"`([^`]+\.(?:jpg|jpeg|png|webp))`", text, re.I)
|
|
if not m:
|
|
m = re.search(r"([\w./\\-]+\.(?:jpg|jpeg|png|webp))", text, re.I)
|
|
if not m:
|
|
return None
|
|
candidate = Path(m.group(1))
|
|
if candidate.is_file():
|
|
return candidate.resolve()
|
|
# Output paths are usually relative to the session dir; resolve via search below.
|
|
return None
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate images with Grok Imagine through the grok CLI"
|
|
)
|
|
parser.add_argument("prompt", help="text description of the image")
|
|
parser.add_argument("-a", "--aspect-ratio", help="e.g. 1:1, 16:9, 9:16, 4:3")
|
|
parser.add_argument("-o", "--output", help="copy the result into this directory")
|
|
args = parser.parse_args()
|
|
|
|
if not shutil.which(GROK):
|
|
die("grok CLI not found. Install with: curl -fsSL https://x.ai/cli/install.sh | bash")
|
|
|
|
full_prompt = f"/imagine {args.prompt}"
|
|
if args.aspect_ratio:
|
|
full_prompt += f", aspect ratio {args.aspect_ratio}"
|
|
|
|
before = time.time()
|
|
proc = subprocess.run(
|
|
[GROK, "-p", full_prompt, "--no-auto-update"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=600,
|
|
)
|
|
output = proc.stdout + proc.stderr
|
|
print(output)
|
|
if proc.returncode != 0:
|
|
die(f"grok exited with code {proc.returncode}")
|
|
|
|
image = newest_image(before - 2) or path_from_output(output)
|
|
if image is None:
|
|
die("could not locate the generated image (check the output above)")
|
|
|
|
image = image.resolve()
|
|
if args.output:
|
|
outdir = Path(args.output)
|
|
outdir.mkdir(parents=True, exist_ok=True)
|
|
dest = outdir / image.name
|
|
if dest.exists():
|
|
dest = outdir / f"{image.stem}-{int(time.time())}{image.suffix}"
|
|
shutil.copy2(image, dest)
|
|
image = dest.resolve()
|
|
print(image)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|