📦 deps(thirdparty): update snapshots
This commit is contained in:
ci[bot]
@@ -1,6 +1,6 @@
|
||||
---
|
||||
name: design
|
||||
description: "Comprehensive design skill: brand identity, design tokens, UI styling, logo generation (55 styles, Gemini AI), corporate identity program (50 deliverables, CIP mockups), HTML presentations (Chart.js), banner design (22 styles, social/ads/web/print), icon design (15 styles, SVG, Gemini 3.1 Pro), social photos (HTML→screenshot, multi-platform). Actions: design logo, create CIP, generate mockups, build slides, design banner, generate icon, create social photos, social media images, brand identity, design system. Platforms: Facebook, Twitter, LinkedIn, YouTube, Instagram, Pinterest, TikTok, Threads, Google Ads."
|
||||
description: "Comprehensive design skill: brand identity, design tokens, UI styling, logo generation (55 styles, Gemini or Atlas Cloud AI), corporate identity program (50 deliverables, CIP mockups), HTML presentations (Chart.js), banner design (22 styles, social/ads/web/print), icon design (15 styles, SVG, Gemini 3.1 Pro), social photos (HTML→screenshot, multi-platform). Actions: design logo, create CIP, generate mockups, build slides, design banner, generate icon, create social photos, social media images, brand identity, design system. Platforms: Facebook, Twitter, LinkedIn, YouTube, Instagram, Pinterest, TikTok, Threads, Google Ads."
|
||||
argument-hint: "[design-type] [context]"
|
||||
license: MIT
|
||||
metadata:
|
||||
@@ -62,6 +62,7 @@ python3 ~/.claude/skills/design/scripts/logo/search.py "healthcare medical" --do
|
||||
```bash
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --brand "TechFlow" --style minimalist --industry tech
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --prompt "coffee shop vintage badge" --style vintage
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --brand "TechFlow" --provider atlas
|
||||
```
|
||||
|
||||
**IMPORTANT:** When scripts fail, try to fix them directly.
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
# Logo Design Reference
|
||||
|
||||
AI-powered logo design with 55+ styles, 30 color palettes, 25 industry guides. Uses Gemini Nano Banana models.
|
||||
AI-powered logo design with 55+ styles, 30 color palettes, 25 industry guides. Gemini Nano Banana is the default provider; Atlas Cloud is also available as an explicit opt-in.
|
||||
|
||||
## Scripts
|
||||
|
||||
| Script | Purpose |
|
||||
|--------|---------|
|
||||
| `scripts/logo/search.py` | Search styles, colors, industries; generate design briefs |
|
||||
| `scripts/logo/generate.py` | Generate logos with Gemini Nano Banana |
|
||||
| `scripts/logo/generate.py` | Generate logos with Gemini Nano Banana or Atlas Cloud |
|
||||
| `scripts/logo/core.py` | BM25 search engine for logo data |
|
||||
|
||||
## Commands
|
||||
@@ -38,9 +38,10 @@ python3 ~/.claude/skills/design/scripts/logo/search.py "healthcare medical" --do
|
||||
```bash
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --brand "TechFlow" --style minimalist --industry tech
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --prompt "coffee shop vintage badge" --style vintage
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --brand "TechFlow" --provider atlas
|
||||
```
|
||||
|
||||
Options: `--style`, `--industry`, `--prompt`
|
||||
Options: `--style`, `--industry`, `--prompt`, `--provider`, `--atlas-model`
|
||||
|
||||
## Available Styles
|
||||
|
||||
@@ -89,4 +90,7 @@ Options: `--style`, `--industry`, `--prompt`
|
||||
```bash
|
||||
export GEMINI_API_KEY="your-key"
|
||||
pip install google-genai
|
||||
|
||||
# Optional Atlas Cloud provider (no extra Python package required)
|
||||
export ATLASCLOUD_API_KEY="your-key"
|
||||
```
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Logo Generation Script using Gemini Nano Banana API
|
||||
Uses Gemini 2.5 Flash Image and Gemini 3 Pro Image Preview models
|
||||
"""Logo generation with Gemini or Atlas Cloud.
|
||||
|
||||
Gemini remains the default provider. Atlas Cloud is opt-in with
|
||||
``--provider atlas`` and uses its asynchronous image generation API.
|
||||
|
||||
Models:
|
||||
- Nano Banana (default): gemini-2.5-flash-image - fast, high-volume, low-latency
|
||||
@@ -13,17 +13,23 @@ Usage:
|
||||
python generate.py --prompt "coffee shop vintage badge" --style vintage --output logo.png
|
||||
python generate.py --brand "TechFlow" --industry tech --style minimalist
|
||||
python generate.py --brand "TechFlow" --pro # Use Nano Banana Pro model
|
||||
python generate.py --brand "TechFlow" --provider atlas
|
||||
|
||||
Batch mode (generates multiple variants):
|
||||
python generate.py --brand "Unikorn" --batch 9 --output-dir ./logos --pro
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from urllib.error import HTTPError, URLError
|
||||
from urllib.parse import urlparse
|
||||
from urllib.request import HTTPRedirectHandler, Request, build_opener
|
||||
|
||||
|
||||
# Load environment variables
|
||||
def load_env():
|
||||
@@ -31,7 +37,7 @@ def load_env():
|
||||
env_paths = [
|
||||
Path(__file__).parent.parent.parent / ".env",
|
||||
Path.home() / ".claude" / "skills" / ".env",
|
||||
Path.home() / ".claude" / ".env"
|
||||
Path.home() / ".claude" / ".env",
|
||||
]
|
||||
|
||||
for env_path in env_paths:
|
||||
@@ -39,29 +45,30 @@ def load_env():
|
||||
with open(env_path) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line and not line.startswith('#') and '=' in line:
|
||||
key, value = line.split('=', 1)
|
||||
if line and not line.startswith("#") and "=" in line:
|
||||
key, value = line.split("=", 1)
|
||||
if key not in os.environ:
|
||||
os.environ[key] = value.strip('"\'')
|
||||
os.environ[key] = value.strip("\"'")
|
||||
|
||||
|
||||
load_env()
|
||||
|
||||
try:
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
except ImportError:
|
||||
print("Error: google-genai package not installed.")
|
||||
print("Install with: pip install google-genai")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
# ============ CONFIGURATION ============
|
||||
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
|
||||
ATLASCLOUD_API_KEY = os.environ.get("ATLASCLOUD_API_KEY")
|
||||
|
||||
# Gemini "Nano Banana" model configurations for image generation
|
||||
GEMINI_FLASH = "gemini-2.5-flash-image" # Nano Banana: fast, high-volume, low-latency
|
||||
GEMINI_PRO = "gemini-3-pro-image-preview" # Nano Banana Pro: professional quality, advanced reasoning
|
||||
|
||||
# Atlas Cloud model validated against the live model catalog and schema.
|
||||
ATLAS_MODEL = "google/nano-banana-2-lite/text-to-image"
|
||||
ATLAS_API_BASE = "https://api.atlascloud.ai/api/v1"
|
||||
HTTP_USER_AGENT = "ui-ux-pro-max/2.5 (Atlas Cloud logo provider)"
|
||||
ATLAS_POLL_INTERVAL = 2
|
||||
ATLAS_MAX_POLLS = 90
|
||||
|
||||
# Supported aspect ratios
|
||||
ASPECT_RATIOS = ["1:1", "16:9", "9:16", "4:3", "3:4"]
|
||||
DEFAULT_ASPECT_RATIO = "1:1" # Square is ideal for logos
|
||||
@@ -99,7 +106,7 @@ STYLE_MODIFIERS = {
|
||||
"mascot": "mascot, character, friendly face, personified, memorable figure",
|
||||
"gradient": "gradient, color transition, vibrant, modern digital feel, smooth color flow",
|
||||
"lineart": "line art, single stroke, continuous line, elegant simplicity, wire-frame style",
|
||||
"negative-space": "negative space, clever use of white space, hidden meaning, dual imagery, optical illusion"
|
||||
"negative-space": "negative space, clever use of white space, hidden meaning, dual imagery, optical illusion",
|
||||
}
|
||||
|
||||
INDUSTRY_PROMPTS = {
|
||||
@@ -112,7 +119,7 @@ INDUSTRY_PROMPTS = {
|
||||
"eco": "eco-friendly, sustainable, natural, green, leaf or earth elements",
|
||||
"education": "education, knowledge, growth, learning, book or cap symbol",
|
||||
"real-estate": "real estate, property, home, roof or building silhouette",
|
||||
"creative": "creative agency, artistic, unique, expressive, colorful"
|
||||
"creative": "creative agency, artistic, unique, expressive, colorful",
|
||||
}
|
||||
|
||||
|
||||
@@ -133,101 +140,274 @@ def enhance_prompt(base_prompt, style=None, industry=None, brand_name=None):
|
||||
return LOGO_PROMPT_TEMPLATE.format(prompt=combined)
|
||||
|
||||
|
||||
def generate_logo(prompt, style=None, industry=None, brand_name=None,
|
||||
output_path=None, use_pro=False, aspect_ratio=None):
|
||||
"""Generate a logo using Gemini models with image generation
|
||||
class _SafeRedirectHandler(HTTPRedirectHandler):
|
||||
"""Reject redirects to non-public or non-HTTPS destinations."""
|
||||
|
||||
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
||||
_validate_public_https_url(newurl)
|
||||
return super().redirect_request(req, fp, code, msg, headers, newurl)
|
||||
|
||||
|
||||
def _validate_public_https_url(url):
|
||||
parsed = urlparse(url)
|
||||
if (
|
||||
parsed.scheme != "https"
|
||||
or not parsed.hostname
|
||||
or parsed.username
|
||||
or parsed.password
|
||||
):
|
||||
raise ValueError("Atlas Cloud returned an invalid media URL")
|
||||
|
||||
hostname = parsed.hostname.lower().rstrip(".")
|
||||
if hostname == "localhost" or hostname.endswith(
|
||||
(".localhost", ".local", ".internal")
|
||||
):
|
||||
raise ValueError("Atlas Cloud media URL used a local hostname")
|
||||
|
||||
try:
|
||||
ip = ipaddress.ip_address(hostname)
|
||||
except ValueError:
|
||||
return
|
||||
else:
|
||||
if not ip.is_global:
|
||||
raise ValueError("Atlas Cloud media URL used a non-public address")
|
||||
|
||||
|
||||
def _json_request(url, api_key, method="GET", payload=None):
|
||||
body = json.dumps(payload).encode("utf-8") if payload is not None else None
|
||||
request = Request(
|
||||
url,
|
||||
data=body,
|
||||
method=method,
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Accept": "application/json",
|
||||
"User-Agent": HTTP_USER_AGENT,
|
||||
**({"Content-Type": "application/json"} if body is not None else {}),
|
||||
},
|
||||
)
|
||||
try:
|
||||
with build_opener(_SafeRedirectHandler()).open(request, timeout=60) as response:
|
||||
return json.loads(response.read().decode("utf-8"))
|
||||
except HTTPError as exc:
|
||||
detail = exc.read().decode("utf-8", errors="replace")
|
||||
raise RuntimeError(
|
||||
f"Atlas Cloud request failed ({exc.code}): {detail[:300]}"
|
||||
) from exc
|
||||
except (URLError, TimeoutError, json.JSONDecodeError) as exc:
|
||||
raise RuntimeError(f"Atlas Cloud request failed: {exc}") from exc
|
||||
|
||||
|
||||
def _atlas_prediction_data(response):
|
||||
if not isinstance(response, dict):
|
||||
raise TypeError("Atlas Cloud returned an invalid response")
|
||||
if response.get("code") not in (None, 0, 200):
|
||||
raise RuntimeError(response.get("message") or "Atlas Cloud request failed")
|
||||
data = response.get("data")
|
||||
if not isinstance(data, dict):
|
||||
raise TypeError("Atlas Cloud response did not include prediction data")
|
||||
return data
|
||||
|
||||
|
||||
def _download_atlas_image(url, output_path):
|
||||
_validate_public_https_url(url)
|
||||
request = Request(
|
||||
url,
|
||||
headers={"Accept": "image/*", "User-Agent": HTTP_USER_AGENT},
|
||||
)
|
||||
try:
|
||||
with build_opener(_SafeRedirectHandler()).open(
|
||||
request, timeout=120
|
||||
) as response:
|
||||
content_type = response.headers.get_content_type()
|
||||
if not content_type.startswith("image/"):
|
||||
raise RuntimeError(
|
||||
f"Atlas Cloud output is not an image ({content_type})"
|
||||
)
|
||||
image_data = response.read()
|
||||
except (HTTPError, URLError, TimeoutError) as exc:
|
||||
raise RuntimeError(f"Unable to download Atlas Cloud image: {exc}") from exc
|
||||
|
||||
if not image_data:
|
||||
raise RuntimeError("Atlas Cloud returned an empty image")
|
||||
with open(output_path, "wb") as output_file:
|
||||
output_file.write(image_data)
|
||||
|
||||
|
||||
def _generate_with_atlas(prompt, output_path, aspect_ratio, api_key, model):
|
||||
if not api_key:
|
||||
raise RuntimeError("ATLASCLOUD_API_KEY not set")
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
}
|
||||
response = _json_request(
|
||||
f"{ATLAS_API_BASE}/model/generateImage",
|
||||
api_key,
|
||||
method="POST",
|
||||
payload=payload,
|
||||
)
|
||||
data = _atlas_prediction_data(response)
|
||||
prediction_id = data.get("id")
|
||||
if not prediction_id:
|
||||
raise RuntimeError("Atlas Cloud did not return a prediction ID")
|
||||
|
||||
for poll_number in range(ATLAS_MAX_POLLS + 1):
|
||||
status = str(data.get("status", "")).lower()
|
||||
if status == "completed":
|
||||
outputs = data.get("outputs")
|
||||
if (
|
||||
not isinstance(outputs, list)
|
||||
or not outputs
|
||||
or not isinstance(outputs[0], str)
|
||||
):
|
||||
raise RuntimeError("Atlas Cloud completed without an image URL")
|
||||
_download_atlas_image(outputs[0], output_path)
|
||||
return
|
||||
if status in {"failed", "timeout", "canceled", "cancelled"}:
|
||||
raise RuntimeError(data.get("error") or f"Atlas Cloud prediction {status}")
|
||||
if poll_number == ATLAS_MAX_POLLS:
|
||||
break
|
||||
time.sleep(ATLAS_POLL_INTERVAL)
|
||||
data = _atlas_prediction_data(
|
||||
_json_request(
|
||||
f"{ATLAS_API_BASE}/model/prediction/{prediction_id}",
|
||||
api_key,
|
||||
)
|
||||
)
|
||||
|
||||
raise RuntimeError("Atlas Cloud prediction timed out while polling")
|
||||
|
||||
|
||||
def _generate_with_gemini(prompt, output_path, aspect_ratio, use_pro):
|
||||
if not GEMINI_API_KEY:
|
||||
raise RuntimeError("GEMINI_API_KEY not set")
|
||||
|
||||
try:
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"google-genai package not installed; run: pip install google-genai"
|
||||
) from exc
|
||||
|
||||
client = genai.Client(api_key=GEMINI_API_KEY)
|
||||
model = GEMINI_PRO if use_pro else GEMINI_FLASH
|
||||
response = client.models.generate_content(
|
||||
model=model,
|
||||
contents=prompt,
|
||||
config=types.GenerateContentConfig(
|
||||
response_modalities=["IMAGE", "TEXT"],
|
||||
image_config=types.ImageConfig(aspect_ratio=aspect_ratio),
|
||||
safety_settings=[
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_HATE_SPEECH",
|
||||
threshold="BLOCK_LOW_AND_ABOVE",
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_DANGEROUS_CONTENT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE",
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE",
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_HARASSMENT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE",
|
||||
),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
for part in response.candidates[0].content.parts:
|
||||
if (
|
||||
hasattr(part, "inline_data")
|
||||
and part.inline_data
|
||||
and part.inline_data.mime_type.startswith("image/")
|
||||
):
|
||||
with open(output_path, "wb") as output_file:
|
||||
output_file.write(part.inline_data.data)
|
||||
return
|
||||
raise RuntimeError("Gemini did not return an image")
|
||||
|
||||
|
||||
def generate_logo(
|
||||
prompt,
|
||||
style=None,
|
||||
industry=None,
|
||||
brand_name=None,
|
||||
output_path=None,
|
||||
use_pro=False,
|
||||
aspect_ratio=None,
|
||||
provider="gemini",
|
||||
atlas_model=ATLAS_MODEL,
|
||||
):
|
||||
"""Generate a logo using Gemini or Atlas Cloud image generation.
|
||||
|
||||
Args:
|
||||
aspect_ratio: Image aspect ratio. Options: "1:1", "16:9", "9:16", "4:3", "3:4"
|
||||
Default is "1:1" (square) for logos.
|
||||
"""
|
||||
|
||||
if not GEMINI_API_KEY:
|
||||
print("Error: GEMINI_API_KEY not set")
|
||||
print("Set it with: export GEMINI_API_KEY='your-key'")
|
||||
return None
|
||||
|
||||
# Initialize client
|
||||
client = genai.Client(api_key=GEMINI_API_KEY)
|
||||
|
||||
# Enhance the prompt
|
||||
full_prompt = enhance_prompt(prompt, style, industry, brand_name)
|
||||
|
||||
# Select model
|
||||
model = GEMINI_PRO if use_pro else GEMINI_FLASH
|
||||
model_label = "Nano Banana Pro (gemini-3-pro-image-preview)" if use_pro else "Nano Banana (gemini-2.5-flash-image)"
|
||||
|
||||
# Set aspect ratio (default to 1:1 for logos)
|
||||
ratio = aspect_ratio if aspect_ratio in ASPECT_RATIOS else DEFAULT_ASPECT_RATIO
|
||||
|
||||
if output_path is None:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") # noqa: DTZ005
|
||||
brand_slug = brand_name.lower().replace(" ", "_") if brand_name else "logo"
|
||||
output_path = f"{brand_slug}_{timestamp}.png"
|
||||
|
||||
if provider == "atlas":
|
||||
model_label = f"Atlas Cloud ({atlas_model})"
|
||||
else:
|
||||
model_label = (
|
||||
"Nano Banana Pro (gemini-3-pro-image-preview)"
|
||||
if use_pro
|
||||
else "Nano Banana (gemini-2.5-flash-image)"
|
||||
)
|
||||
|
||||
print(f"Generating logo with {model_label}...")
|
||||
print(f"Aspect ratio: {ratio}")
|
||||
print(f"Prompt: {full_prompt[:150]}...")
|
||||
print()
|
||||
|
||||
try:
|
||||
# Generate image using Gemini with image generation capability
|
||||
response = client.models.generate_content(
|
||||
model=model,
|
||||
contents=full_prompt,
|
||||
config=types.GenerateContentConfig(
|
||||
response_modalities=["IMAGE", "TEXT"],
|
||||
image_config=types.ImageConfig(
|
||||
aspect_ratio=ratio
|
||||
),
|
||||
safety_settings=[
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_HATE_SPEECH",
|
||||
threshold="BLOCK_LOW_AND_ABOVE"
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_DANGEROUS_CONTENT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE"
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE"
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_HARASSMENT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE"
|
||||
),
|
||||
]
|
||||
if provider == "atlas":
|
||||
_generate_with_atlas(
|
||||
full_prompt,
|
||||
output_path,
|
||||
ratio,
|
||||
ATLASCLOUD_API_KEY,
|
||||
atlas_model,
|
||||
)
|
||||
)
|
||||
|
||||
# Extract image from response
|
||||
image_data = None
|
||||
for part in response.candidates[0].content.parts:
|
||||
if hasattr(part, 'inline_data') and part.inline_data:
|
||||
if part.inline_data.mime_type.startswith('image/'):
|
||||
image_data = part.inline_data.data
|
||||
break
|
||||
|
||||
if not image_data:
|
||||
print("No image generated. The model may not have produced an image.")
|
||||
print("Try a different prompt or check if the model supports image generation.")
|
||||
return None
|
||||
|
||||
# Determine output path
|
||||
if output_path is None:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
brand_slug = brand_name.lower().replace(" ", "_") if brand_name else "logo"
|
||||
output_path = f"{brand_slug}_{timestamp}.png"
|
||||
|
||||
# Save image
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(image_data)
|
||||
else:
|
||||
_generate_with_gemini(full_prompt, output_path, ratio, use_pro)
|
||||
|
||||
print(f"Logo saved to: {output_path}")
|
||||
return output_path
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error generating logo: {e}")
|
||||
except Exception as exc: # noqa: BLE001 - provider SDK errors are not standardized
|
||||
print(f"Error generating logo: {exc}")
|
||||
return None
|
||||
|
||||
|
||||
def generate_batch(prompt, brand_name, count, output_dir, use_pro=False, brand_context=None, aspect_ratio=None):
|
||||
def generate_batch(
|
||||
prompt,
|
||||
brand_name,
|
||||
count,
|
||||
output_dir,
|
||||
use_pro=False,
|
||||
brand_context=None,
|
||||
aspect_ratio=None,
|
||||
provider="gemini",
|
||||
atlas_model=ATLAS_MODEL,
|
||||
):
|
||||
"""Generate multiple logo variants with different styles"""
|
||||
|
||||
# Select appropriate styles for batch generation
|
||||
@@ -247,16 +427,20 @@ def generate_batch(prompt, brand_name, count, output_dir, use_pro=False, brand_c
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
results = []
|
||||
model_label = "Pro" if use_pro else "Flash"
|
||||
model_label = (
|
||||
f"Atlas Cloud ({atlas_model})"
|
||||
if provider == "atlas"
|
||||
else f"Nano Banana {'Pro' if use_pro else 'Flash'}"
|
||||
)
|
||||
ratio = aspect_ratio if aspect_ratio in ASPECT_RATIOS else DEFAULT_ASPECT_RATIO
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" BATCH LOGO GENERATION: {brand_name}")
|
||||
print(f" Model: Nano Banana {model_label}")
|
||||
print(f" Model: {model_label}")
|
||||
print(f" Aspect Ratio: {ratio}")
|
||||
print(f" Variants: {count}")
|
||||
print(f" Output: {output_dir}")
|
||||
print(f"{'='*60}\n")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
for i in range(min(count, len(batch_styles))):
|
||||
style_key, style_desc = batch_styles[i]
|
||||
@@ -267,10 +451,10 @@ def generate_batch(prompt, brand_name, count, output_dir, use_pro=False, brand_c
|
||||
enhanced_prompt = f"{brand_context}, {enhanced_prompt}"
|
||||
|
||||
# Generate filename
|
||||
filename = f"{brand_name.lower().replace(' ', '_')}_{style_key}_{i+1:02d}.png"
|
||||
filename = f"{brand_name.lower().replace(' ', '_')}_{style_key}_{i + 1:02d}.png"
|
||||
output_path = os.path.join(output_dir, filename)
|
||||
|
||||
print(f"[{i+1}/{count}] Generating {style_key} variant...")
|
||||
print(f"[{i + 1}/{count}] Generating {style_key} variant...")
|
||||
|
||||
result = generate_logo(
|
||||
prompt=enhanced_prompt,
|
||||
@@ -279,7 +463,9 @@ def generate_batch(prompt, brand_name, count, output_dir, use_pro=False, brand_c
|
||||
brand_name=brand_name,
|
||||
output_path=output_path,
|
||||
use_pro=use_pro,
|
||||
aspect_ratio=aspect_ratio
|
||||
aspect_ratio=aspect_ratio,
|
||||
provider=provider,
|
||||
atlas_model=atlas_model,
|
||||
)
|
||||
|
||||
if result:
|
||||
@@ -292,31 +478,70 @@ def generate_batch(prompt, brand_name, count, output_dir, use_pro=False, brand_c
|
||||
if i < count - 1:
|
||||
time.sleep(2)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" BATCH COMPLETE: {len(results)}/{count} logos generated")
|
||||
print(f"{'='*60}\n")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate logos using Gemini Nano Banana models")
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate logos using Gemini or Atlas Cloud"
|
||||
)
|
||||
parser.add_argument("--prompt", "-p", type=str, help="Logo description prompt")
|
||||
parser.add_argument("--brand", "-b", type=str, help="Brand name")
|
||||
parser.add_argument("--style", "-s", choices=list(STYLE_MODIFIERS.keys()), help="Logo style")
|
||||
parser.add_argument("--industry", "-i", choices=list(INDUSTRY_PROMPTS.keys()), help="Industry type")
|
||||
parser.add_argument(
|
||||
"--style", "-s", choices=list(STYLE_MODIFIERS.keys()), help="Logo style"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--industry", "-i", choices=list(INDUSTRY_PROMPTS.keys()), help="Industry type"
|
||||
)
|
||||
parser.add_argument("--output", "-o", type=str, help="Output file path")
|
||||
parser.add_argument("--output-dir", type=str, help="Output directory for batch generation")
|
||||
parser.add_argument("--batch", type=int, help="Number of logo variants to generate (batch mode)")
|
||||
parser.add_argument("--brand-context", type=str, help="Additional brand context for prompts")
|
||||
parser.add_argument("--pro", action="store_true", help="Use Nano Banana Pro (gemini-3-pro-image-preview) for professional quality")
|
||||
parser.add_argument("--aspect-ratio", "-r", choices=ASPECT_RATIOS, default=DEFAULT_ASPECT_RATIO,
|
||||
help=f"Image aspect ratio (default: {DEFAULT_ASPECT_RATIO} for logos)")
|
||||
parser.add_argument("--list-styles", action="store_true", help="List available styles")
|
||||
parser.add_argument("--list-industries", action="store_true", help="List available industries")
|
||||
parser.add_argument(
|
||||
"--output-dir", type=str, help="Output directory for batch generation"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch", type=int, help="Number of logo variants to generate (batch mode)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--brand-context", type=str, help="Additional brand context for prompts"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pro",
|
||||
action="store_true",
|
||||
help="Use Nano Banana Pro (gemini-3-pro-image-preview) for professional quality",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--provider",
|
||||
choices=["gemini", "atlas"],
|
||||
default="gemini",
|
||||
help="Image provider (default: gemini)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--atlas-model",
|
||||
default=ATLAS_MODEL,
|
||||
help=f"Atlas Cloud image model (default: {ATLAS_MODEL})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--aspect-ratio",
|
||||
"-r",
|
||||
choices=ASPECT_RATIOS,
|
||||
default=DEFAULT_ASPECT_RATIO,
|
||||
help=f"Image aspect ratio (default: {DEFAULT_ASPECT_RATIO} for logos)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--list-styles", action="store_true", help="List available styles"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--list-industries", action="store_true", help="List available industries"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.provider == "atlas" and args.pro:
|
||||
parser.error("--pro is only available with --provider gemini")
|
||||
|
||||
if args.list_styles:
|
||||
print("Available styles:")
|
||||
for style, desc in STYLE_MODIFIERS.items():
|
||||
@@ -336,7 +561,9 @@ def main():
|
||||
|
||||
# Batch mode
|
||||
if args.batch:
|
||||
output_dir = args.output_dir or f"./{args.brand.lower().replace(' ', '_')}_logos"
|
||||
output_dir = (
|
||||
args.output_dir or f"./{args.brand.lower().replace(' ', '_')}_logos"
|
||||
)
|
||||
generate_batch(
|
||||
prompt=prompt,
|
||||
brand_name=args.brand or "Logo",
|
||||
@@ -344,7 +571,9 @@ def main():
|
||||
output_dir=output_dir,
|
||||
use_pro=args.pro,
|
||||
brand_context=args.brand_context,
|
||||
aspect_ratio=args.aspect_ratio
|
||||
aspect_ratio=args.aspect_ratio,
|
||||
provider=args.provider,
|
||||
atlas_model=args.atlas_model,
|
||||
)
|
||||
else:
|
||||
generate_logo(
|
||||
@@ -354,7 +583,9 @@ def main():
|
||||
brand_name=args.brand,
|
||||
output_path=args.output,
|
||||
use_pro=args.pro,
|
||||
aspect_ratio=args.aspect_ratio
|
||||
aspect_ratio=args.aspect_ratio,
|
||||
provider=args.provider,
|
||||
atlas_model=args.atlas_model,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
import importlib.util
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import call, patch
|
||||
|
||||
MODULE_PATH = Path(__file__).parents[1] / "generate.py"
|
||||
SPEC = importlib.util.spec_from_file_location("logo_generate", MODULE_PATH)
|
||||
logo_generate = importlib.util.module_from_spec(SPEC)
|
||||
SPEC.loader.exec_module(logo_generate)
|
||||
|
||||
|
||||
class AtlasGenerationTests(unittest.TestCase):
|
||||
@patch.object(logo_generate, "_download_atlas_image")
|
||||
@patch.object(logo_generate.time, "sleep")
|
||||
@patch.object(logo_generate, "_json_request")
|
||||
def test_atlas_submits_once_and_polls_until_completed(
|
||||
self, json_request, sleep, download
|
||||
):
|
||||
json_request.side_effect = [
|
||||
{"code": 200, "data": {"id": "pred-123", "status": "created"}},
|
||||
{"code": 200, "data": {"id": "pred-123", "status": "processing"}},
|
||||
{
|
||||
"code": 200,
|
||||
"data": {
|
||||
"id": "pred-123",
|
||||
"status": "completed",
|
||||
"outputs": ["https://media.example.com/logo.png"],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
logo_generate._generate_with_atlas(
|
||||
"logo prompt", "logo.png", "1:1", "atlas-key", "atlas/model"
|
||||
)
|
||||
|
||||
self.assertEqual(json_request.call_count, 3)
|
||||
self.assertEqual(
|
||||
json_request.call_args_list[0],
|
||||
call(
|
||||
f"{logo_generate.ATLAS_API_BASE}/model/generateImage",
|
||||
"atlas-key",
|
||||
method="POST",
|
||||
payload={
|
||||
"model": "atlas/model",
|
||||
"prompt": "logo prompt",
|
||||
"aspect_ratio": "1:1",
|
||||
},
|
||||
),
|
||||
)
|
||||
self.assertEqual(
|
||||
json_request.call_args_list[1:],
|
||||
[
|
||||
call(
|
||||
f"{logo_generate.ATLAS_API_BASE}/model/prediction/pred-123",
|
||||
"atlas-key",
|
||||
),
|
||||
call(
|
||||
f"{logo_generate.ATLAS_API_BASE}/model/prediction/pred-123",
|
||||
"atlas-key",
|
||||
),
|
||||
],
|
||||
)
|
||||
self.assertEqual(sleep.call_count, 2)
|
||||
download.assert_called_once_with(
|
||||
"https://media.example.com/logo.png", "logo.png"
|
||||
)
|
||||
|
||||
@patch.object(logo_generate, "_json_request")
|
||||
def test_atlas_does_not_retry_generation_post(self, json_request):
|
||||
json_request.side_effect = RuntimeError("network error")
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "network error"):
|
||||
logo_generate._generate_with_atlas(
|
||||
"logo prompt", "logo.png", "1:1", "atlas-key", "atlas/model"
|
||||
)
|
||||
|
||||
json_request.assert_called_once()
|
||||
|
||||
@patch.object(logo_generate, "_validate_public_https_url")
|
||||
@patch.object(logo_generate, "build_opener")
|
||||
def test_media_download_never_forwards_api_key(self, build_opener, validate):
|
||||
class Headers:
|
||||
@staticmethod
|
||||
def get_content_type():
|
||||
return "image/png"
|
||||
|
||||
class Response:
|
||||
headers = Headers()
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def read():
|
||||
return b"png-bytes"
|
||||
|
||||
build_opener.return_value.open.return_value = Response()
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
output = Path(temp_dir) / "logo.png"
|
||||
logo_generate._download_atlas_image(
|
||||
"https://media.example.com/logo.png", output
|
||||
)
|
||||
self.assertEqual(output.read_bytes(), b"png-bytes")
|
||||
|
||||
request = build_opener.return_value.open.call_args.args[0]
|
||||
headers = {key.lower(): value for key, value in request.header_items()}
|
||||
self.assertNotIn("authorization", headers)
|
||||
self.assertEqual(headers["accept"], "image/*")
|
||||
self.assertEqual(headers["user-agent"], logo_generate.HTTP_USER_AGENT)
|
||||
validate.assert_called_once_with("https://media.example.com/logo.png")
|
||||
|
||||
def test_atlas_requires_api_key(self):
|
||||
with self.assertRaisesRegex(RuntimeError, "ATLASCLOUD_API_KEY not set"):
|
||||
logo_generate._generate_with_atlas(
|
||||
"logo prompt", "logo.png", "1:1", None, "atlas/model"
|
||||
)
|
||||
|
||||
def test_media_url_rejects_private_addresses(self):
|
||||
with self.assertRaisesRegex(ValueError, "non-public address"):
|
||||
logo_generate._validate_public_https_url("https://127.0.0.1/logo.png")
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "local hostname"):
|
||||
logo_generate._validate_public_https_url("https://assets.local/logo.png")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,8 +1,8 @@
|
||||
# Source
|
||||
|
||||
- Repo: https://github.com/nextlevelbuilder/ui-ux-pro-max-skill
|
||||
- Ref: e4f45473691e4b389519ee4bc359a3d6df666c26
|
||||
- Ref: 8bd29e775453ebcae52b6e6514fbf134df0c5770
|
||||
- Remove-Paths:
|
||||
- Snapshot: 2026-08-26
|
||||
- Snapshot: 2026-08-27
|
||||
- Sync-Mode: render_skill
|
||||
- Notes: vendored into playbook branch thirdparty/skill
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
---
|
||||
name: design
|
||||
description: "Comprehensive design skill: brand identity, design tokens, UI styling, logo generation (55 styles, Gemini AI), corporate identity program (50 deliverables, CIP mockups), HTML presentations (Chart.js), banner design (22 styles, social/ads/web/print), icon design (15 styles, SVG, Gemini 3.1 Pro), social photos (HTML→screenshot, multi-platform). Actions: design logo, create CIP, generate mockups, build slides, design banner, generate icon, create social photos, social media images, brand identity, design system. Platforms: Facebook, Twitter, LinkedIn, YouTube, Instagram, Pinterest, TikTok, Threads, Google Ads."
|
||||
description: "Comprehensive design skill: brand identity, design tokens, UI styling, logo generation (55 styles, Gemini or Atlas Cloud AI), corporate identity program (50 deliverables, CIP mockups), HTML presentations (Chart.js), banner design (22 styles, social/ads/web/print), icon design (15 styles, SVG, Gemini 3.1 Pro), social photos (HTML→screenshot, multi-platform). Actions: design logo, create CIP, generate mockups, build slides, design banner, generate icon, create social photos, social media images, brand identity, design system. Platforms: Facebook, Twitter, LinkedIn, YouTube, Instagram, Pinterest, TikTok, Threads, Google Ads."
|
||||
argument-hint: "[design-type] [context]"
|
||||
license: MIT
|
||||
metadata:
|
||||
@@ -62,6 +62,7 @@ python3 ~/.claude/skills/design/scripts/logo/search.py "healthcare medical" --do
|
||||
```bash
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --brand "TechFlow" --style minimalist --industry tech
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --prompt "coffee shop vintage badge" --style vintage
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --brand "TechFlow" --provider atlas
|
||||
```
|
||||
|
||||
**IMPORTANT:** When scripts fail, try to fix them directly.
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
# Logo Design Reference
|
||||
|
||||
AI-powered logo design with 55+ styles, 30 color palettes, 25 industry guides. Uses Gemini Nano Banana models.
|
||||
AI-powered logo design with 55+ styles, 30 color palettes, 25 industry guides. Gemini Nano Banana is the default provider; Atlas Cloud is also available as an explicit opt-in.
|
||||
|
||||
## Scripts
|
||||
|
||||
| Script | Purpose |
|
||||
|--------|---------|
|
||||
| `scripts/logo/search.py` | Search styles, colors, industries; generate design briefs |
|
||||
| `scripts/logo/generate.py` | Generate logos with Gemini Nano Banana |
|
||||
| `scripts/logo/generate.py` | Generate logos with Gemini Nano Banana or Atlas Cloud |
|
||||
| `scripts/logo/core.py` | BM25 search engine for logo data |
|
||||
|
||||
## Commands
|
||||
@@ -38,9 +38,10 @@ python3 ~/.claude/skills/design/scripts/logo/search.py "healthcare medical" --do
|
||||
```bash
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --brand "TechFlow" --style minimalist --industry tech
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --prompt "coffee shop vintage badge" --style vintage
|
||||
python3 ~/.claude/skills/design/scripts/logo/generate.py --brand "TechFlow" --provider atlas
|
||||
```
|
||||
|
||||
Options: `--style`, `--industry`, `--prompt`
|
||||
Options: `--style`, `--industry`, `--prompt`, `--provider`, `--atlas-model`
|
||||
|
||||
## Available Styles
|
||||
|
||||
@@ -89,4 +90,7 @@ Options: `--style`, `--industry`, `--prompt`
|
||||
```bash
|
||||
export GEMINI_API_KEY="your-key"
|
||||
pip install google-genai
|
||||
|
||||
# Optional Atlas Cloud provider (no extra Python package required)
|
||||
export ATLASCLOUD_API_KEY="your-key"
|
||||
```
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Logo Generation Script using Gemini Nano Banana API
|
||||
Uses Gemini 2.5 Flash Image and Gemini 3 Pro Image Preview models
|
||||
"""Logo generation with Gemini or Atlas Cloud.
|
||||
|
||||
Gemini remains the default provider. Atlas Cloud is opt-in with
|
||||
``--provider atlas`` and uses its asynchronous image generation API.
|
||||
|
||||
Models:
|
||||
- Nano Banana (default): gemini-2.5-flash-image - fast, high-volume, low-latency
|
||||
@@ -13,17 +13,23 @@ Usage:
|
||||
python generate.py --prompt "coffee shop vintage badge" --style vintage --output logo.png
|
||||
python generate.py --brand "TechFlow" --industry tech --style minimalist
|
||||
python generate.py --brand "TechFlow" --pro # Use Nano Banana Pro model
|
||||
python generate.py --brand "TechFlow" --provider atlas
|
||||
|
||||
Batch mode (generates multiple variants):
|
||||
python generate.py --brand "Unikorn" --batch 9 --output-dir ./logos --pro
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from urllib.error import HTTPError, URLError
|
||||
from urllib.parse import urlparse
|
||||
from urllib.request import HTTPRedirectHandler, Request, build_opener
|
||||
|
||||
|
||||
# Load environment variables
|
||||
def load_env():
|
||||
@@ -31,7 +37,7 @@ def load_env():
|
||||
env_paths = [
|
||||
Path(__file__).parent.parent.parent / ".env",
|
||||
Path.home() / ".claude" / "skills" / ".env",
|
||||
Path.home() / ".claude" / ".env"
|
||||
Path.home() / ".claude" / ".env",
|
||||
]
|
||||
|
||||
for env_path in env_paths:
|
||||
@@ -39,29 +45,30 @@ def load_env():
|
||||
with open(env_path) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line and not line.startswith('#') and '=' in line:
|
||||
key, value = line.split('=', 1)
|
||||
if line and not line.startswith("#") and "=" in line:
|
||||
key, value = line.split("=", 1)
|
||||
if key not in os.environ:
|
||||
os.environ[key] = value.strip('"\'')
|
||||
os.environ[key] = value.strip("\"'")
|
||||
|
||||
|
||||
load_env()
|
||||
|
||||
try:
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
except ImportError:
|
||||
print("Error: google-genai package not installed.")
|
||||
print("Install with: pip install google-genai")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
# ============ CONFIGURATION ============
|
||||
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
|
||||
ATLASCLOUD_API_KEY = os.environ.get("ATLASCLOUD_API_KEY")
|
||||
|
||||
# Gemini "Nano Banana" model configurations for image generation
|
||||
GEMINI_FLASH = "gemini-2.5-flash-image" # Nano Banana: fast, high-volume, low-latency
|
||||
GEMINI_PRO = "gemini-3-pro-image-preview" # Nano Banana Pro: professional quality, advanced reasoning
|
||||
|
||||
# Atlas Cloud model validated against the live model catalog and schema.
|
||||
ATLAS_MODEL = "google/nano-banana-2-lite/text-to-image"
|
||||
ATLAS_API_BASE = "https://api.atlascloud.ai/api/v1"
|
||||
HTTP_USER_AGENT = "ui-ux-pro-max/2.5 (Atlas Cloud logo provider)"
|
||||
ATLAS_POLL_INTERVAL = 2
|
||||
ATLAS_MAX_POLLS = 90
|
||||
|
||||
# Supported aspect ratios
|
||||
ASPECT_RATIOS = ["1:1", "16:9", "9:16", "4:3", "3:4"]
|
||||
DEFAULT_ASPECT_RATIO = "1:1" # Square is ideal for logos
|
||||
@@ -99,7 +106,7 @@ STYLE_MODIFIERS = {
|
||||
"mascot": "mascot, character, friendly face, personified, memorable figure",
|
||||
"gradient": "gradient, color transition, vibrant, modern digital feel, smooth color flow",
|
||||
"lineart": "line art, single stroke, continuous line, elegant simplicity, wire-frame style",
|
||||
"negative-space": "negative space, clever use of white space, hidden meaning, dual imagery, optical illusion"
|
||||
"negative-space": "negative space, clever use of white space, hidden meaning, dual imagery, optical illusion",
|
||||
}
|
||||
|
||||
INDUSTRY_PROMPTS = {
|
||||
@@ -112,7 +119,7 @@ INDUSTRY_PROMPTS = {
|
||||
"eco": "eco-friendly, sustainable, natural, green, leaf or earth elements",
|
||||
"education": "education, knowledge, growth, learning, book or cap symbol",
|
||||
"real-estate": "real estate, property, home, roof or building silhouette",
|
||||
"creative": "creative agency, artistic, unique, expressive, colorful"
|
||||
"creative": "creative agency, artistic, unique, expressive, colorful",
|
||||
}
|
||||
|
||||
|
||||
@@ -133,101 +140,274 @@ def enhance_prompt(base_prompt, style=None, industry=None, brand_name=None):
|
||||
return LOGO_PROMPT_TEMPLATE.format(prompt=combined)
|
||||
|
||||
|
||||
def generate_logo(prompt, style=None, industry=None, brand_name=None,
|
||||
output_path=None, use_pro=False, aspect_ratio=None):
|
||||
"""Generate a logo using Gemini models with image generation
|
||||
class _SafeRedirectHandler(HTTPRedirectHandler):
|
||||
"""Reject redirects to non-public or non-HTTPS destinations."""
|
||||
|
||||
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
||||
_validate_public_https_url(newurl)
|
||||
return super().redirect_request(req, fp, code, msg, headers, newurl)
|
||||
|
||||
|
||||
def _validate_public_https_url(url):
|
||||
parsed = urlparse(url)
|
||||
if (
|
||||
parsed.scheme != "https"
|
||||
or not parsed.hostname
|
||||
or parsed.username
|
||||
or parsed.password
|
||||
):
|
||||
raise ValueError("Atlas Cloud returned an invalid media URL")
|
||||
|
||||
hostname = parsed.hostname.lower().rstrip(".")
|
||||
if hostname == "localhost" or hostname.endswith(
|
||||
(".localhost", ".local", ".internal")
|
||||
):
|
||||
raise ValueError("Atlas Cloud media URL used a local hostname")
|
||||
|
||||
try:
|
||||
ip = ipaddress.ip_address(hostname)
|
||||
except ValueError:
|
||||
return
|
||||
else:
|
||||
if not ip.is_global:
|
||||
raise ValueError("Atlas Cloud media URL used a non-public address")
|
||||
|
||||
|
||||
def _json_request(url, api_key, method="GET", payload=None):
|
||||
body = json.dumps(payload).encode("utf-8") if payload is not None else None
|
||||
request = Request(
|
||||
url,
|
||||
data=body,
|
||||
method=method,
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Accept": "application/json",
|
||||
"User-Agent": HTTP_USER_AGENT,
|
||||
**({"Content-Type": "application/json"} if body is not None else {}),
|
||||
},
|
||||
)
|
||||
try:
|
||||
with build_opener(_SafeRedirectHandler()).open(request, timeout=60) as response:
|
||||
return json.loads(response.read().decode("utf-8"))
|
||||
except HTTPError as exc:
|
||||
detail = exc.read().decode("utf-8", errors="replace")
|
||||
raise RuntimeError(
|
||||
f"Atlas Cloud request failed ({exc.code}): {detail[:300]}"
|
||||
) from exc
|
||||
except (URLError, TimeoutError, json.JSONDecodeError) as exc:
|
||||
raise RuntimeError(f"Atlas Cloud request failed: {exc}") from exc
|
||||
|
||||
|
||||
def _atlas_prediction_data(response):
|
||||
if not isinstance(response, dict):
|
||||
raise TypeError("Atlas Cloud returned an invalid response")
|
||||
if response.get("code") not in (None, 0, 200):
|
||||
raise RuntimeError(response.get("message") or "Atlas Cloud request failed")
|
||||
data = response.get("data")
|
||||
if not isinstance(data, dict):
|
||||
raise TypeError("Atlas Cloud response did not include prediction data")
|
||||
return data
|
||||
|
||||
|
||||
def _download_atlas_image(url, output_path):
|
||||
_validate_public_https_url(url)
|
||||
request = Request(
|
||||
url,
|
||||
headers={"Accept": "image/*", "User-Agent": HTTP_USER_AGENT},
|
||||
)
|
||||
try:
|
||||
with build_opener(_SafeRedirectHandler()).open(
|
||||
request, timeout=120
|
||||
) as response:
|
||||
content_type = response.headers.get_content_type()
|
||||
if not content_type.startswith("image/"):
|
||||
raise RuntimeError(
|
||||
f"Atlas Cloud output is not an image ({content_type})"
|
||||
)
|
||||
image_data = response.read()
|
||||
except (HTTPError, URLError, TimeoutError) as exc:
|
||||
raise RuntimeError(f"Unable to download Atlas Cloud image: {exc}") from exc
|
||||
|
||||
if not image_data:
|
||||
raise RuntimeError("Atlas Cloud returned an empty image")
|
||||
with open(output_path, "wb") as output_file:
|
||||
output_file.write(image_data)
|
||||
|
||||
|
||||
def _generate_with_atlas(prompt, output_path, aspect_ratio, api_key, model):
|
||||
if not api_key:
|
||||
raise RuntimeError("ATLASCLOUD_API_KEY not set")
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
}
|
||||
response = _json_request(
|
||||
f"{ATLAS_API_BASE}/model/generateImage",
|
||||
api_key,
|
||||
method="POST",
|
||||
payload=payload,
|
||||
)
|
||||
data = _atlas_prediction_data(response)
|
||||
prediction_id = data.get("id")
|
||||
if not prediction_id:
|
||||
raise RuntimeError("Atlas Cloud did not return a prediction ID")
|
||||
|
||||
for poll_number in range(ATLAS_MAX_POLLS + 1):
|
||||
status = str(data.get("status", "")).lower()
|
||||
if status == "completed":
|
||||
outputs = data.get("outputs")
|
||||
if (
|
||||
not isinstance(outputs, list)
|
||||
or not outputs
|
||||
or not isinstance(outputs[0], str)
|
||||
):
|
||||
raise RuntimeError("Atlas Cloud completed without an image URL")
|
||||
_download_atlas_image(outputs[0], output_path)
|
||||
return
|
||||
if status in {"failed", "timeout", "canceled", "cancelled"}:
|
||||
raise RuntimeError(data.get("error") or f"Atlas Cloud prediction {status}")
|
||||
if poll_number == ATLAS_MAX_POLLS:
|
||||
break
|
||||
time.sleep(ATLAS_POLL_INTERVAL)
|
||||
data = _atlas_prediction_data(
|
||||
_json_request(
|
||||
f"{ATLAS_API_BASE}/model/prediction/{prediction_id}",
|
||||
api_key,
|
||||
)
|
||||
)
|
||||
|
||||
raise RuntimeError("Atlas Cloud prediction timed out while polling")
|
||||
|
||||
|
||||
def _generate_with_gemini(prompt, output_path, aspect_ratio, use_pro):
|
||||
if not GEMINI_API_KEY:
|
||||
raise RuntimeError("GEMINI_API_KEY not set")
|
||||
|
||||
try:
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"google-genai package not installed; run: pip install google-genai"
|
||||
) from exc
|
||||
|
||||
client = genai.Client(api_key=GEMINI_API_KEY)
|
||||
model = GEMINI_PRO if use_pro else GEMINI_FLASH
|
||||
response = client.models.generate_content(
|
||||
model=model,
|
||||
contents=prompt,
|
||||
config=types.GenerateContentConfig(
|
||||
response_modalities=["IMAGE", "TEXT"],
|
||||
image_config=types.ImageConfig(aspect_ratio=aspect_ratio),
|
||||
safety_settings=[
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_HATE_SPEECH",
|
||||
threshold="BLOCK_LOW_AND_ABOVE",
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_DANGEROUS_CONTENT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE",
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE",
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_HARASSMENT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE",
|
||||
),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
for part in response.candidates[0].content.parts:
|
||||
if (
|
||||
hasattr(part, "inline_data")
|
||||
and part.inline_data
|
||||
and part.inline_data.mime_type.startswith("image/")
|
||||
):
|
||||
with open(output_path, "wb") as output_file:
|
||||
output_file.write(part.inline_data.data)
|
||||
return
|
||||
raise RuntimeError("Gemini did not return an image")
|
||||
|
||||
|
||||
def generate_logo(
|
||||
prompt,
|
||||
style=None,
|
||||
industry=None,
|
||||
brand_name=None,
|
||||
output_path=None,
|
||||
use_pro=False,
|
||||
aspect_ratio=None,
|
||||
provider="gemini",
|
||||
atlas_model=ATLAS_MODEL,
|
||||
):
|
||||
"""Generate a logo using Gemini or Atlas Cloud image generation.
|
||||
|
||||
Args:
|
||||
aspect_ratio: Image aspect ratio. Options: "1:1", "16:9", "9:16", "4:3", "3:4"
|
||||
Default is "1:1" (square) for logos.
|
||||
"""
|
||||
|
||||
if not GEMINI_API_KEY:
|
||||
print("Error: GEMINI_API_KEY not set")
|
||||
print("Set it with: export GEMINI_API_KEY='your-key'")
|
||||
return None
|
||||
|
||||
# Initialize client
|
||||
client = genai.Client(api_key=GEMINI_API_KEY)
|
||||
|
||||
# Enhance the prompt
|
||||
full_prompt = enhance_prompt(prompt, style, industry, brand_name)
|
||||
|
||||
# Select model
|
||||
model = GEMINI_PRO if use_pro else GEMINI_FLASH
|
||||
model_label = "Nano Banana Pro (gemini-3-pro-image-preview)" if use_pro else "Nano Banana (gemini-2.5-flash-image)"
|
||||
|
||||
# Set aspect ratio (default to 1:1 for logos)
|
||||
ratio = aspect_ratio if aspect_ratio in ASPECT_RATIOS else DEFAULT_ASPECT_RATIO
|
||||
|
||||
if output_path is None:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") # noqa: DTZ005
|
||||
brand_slug = brand_name.lower().replace(" ", "_") if brand_name else "logo"
|
||||
output_path = f"{brand_slug}_{timestamp}.png"
|
||||
|
||||
if provider == "atlas":
|
||||
model_label = f"Atlas Cloud ({atlas_model})"
|
||||
else:
|
||||
model_label = (
|
||||
"Nano Banana Pro (gemini-3-pro-image-preview)"
|
||||
if use_pro
|
||||
else "Nano Banana (gemini-2.5-flash-image)"
|
||||
)
|
||||
|
||||
print(f"Generating logo with {model_label}...")
|
||||
print(f"Aspect ratio: {ratio}")
|
||||
print(f"Prompt: {full_prompt[:150]}...")
|
||||
print()
|
||||
|
||||
try:
|
||||
# Generate image using Gemini with image generation capability
|
||||
response = client.models.generate_content(
|
||||
model=model,
|
||||
contents=full_prompt,
|
||||
config=types.GenerateContentConfig(
|
||||
response_modalities=["IMAGE", "TEXT"],
|
||||
image_config=types.ImageConfig(
|
||||
aspect_ratio=ratio
|
||||
),
|
||||
safety_settings=[
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_HATE_SPEECH",
|
||||
threshold="BLOCK_LOW_AND_ABOVE"
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_DANGEROUS_CONTENT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE"
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE"
|
||||
),
|
||||
types.SafetySetting(
|
||||
category="HARM_CATEGORY_HARASSMENT",
|
||||
threshold="BLOCK_LOW_AND_ABOVE"
|
||||
),
|
||||
]
|
||||
if provider == "atlas":
|
||||
_generate_with_atlas(
|
||||
full_prompt,
|
||||
output_path,
|
||||
ratio,
|
||||
ATLASCLOUD_API_KEY,
|
||||
atlas_model,
|
||||
)
|
||||
)
|
||||
|
||||
# Extract image from response
|
||||
image_data = None
|
||||
for part in response.candidates[0].content.parts:
|
||||
if hasattr(part, 'inline_data') and part.inline_data:
|
||||
if part.inline_data.mime_type.startswith('image/'):
|
||||
image_data = part.inline_data.data
|
||||
break
|
||||
|
||||
if not image_data:
|
||||
print("No image generated. The model may not have produced an image.")
|
||||
print("Try a different prompt or check if the model supports image generation.")
|
||||
return None
|
||||
|
||||
# Determine output path
|
||||
if output_path is None:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
brand_slug = brand_name.lower().replace(" ", "_") if brand_name else "logo"
|
||||
output_path = f"{brand_slug}_{timestamp}.png"
|
||||
|
||||
# Save image
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(image_data)
|
||||
else:
|
||||
_generate_with_gemini(full_prompt, output_path, ratio, use_pro)
|
||||
|
||||
print(f"Logo saved to: {output_path}")
|
||||
return output_path
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error generating logo: {e}")
|
||||
except Exception as exc: # noqa: BLE001 - provider SDK errors are not standardized
|
||||
print(f"Error generating logo: {exc}")
|
||||
return None
|
||||
|
||||
|
||||
def generate_batch(prompt, brand_name, count, output_dir, use_pro=False, brand_context=None, aspect_ratio=None):
|
||||
def generate_batch(
|
||||
prompt,
|
||||
brand_name,
|
||||
count,
|
||||
output_dir,
|
||||
use_pro=False,
|
||||
brand_context=None,
|
||||
aspect_ratio=None,
|
||||
provider="gemini",
|
||||
atlas_model=ATLAS_MODEL,
|
||||
):
|
||||
"""Generate multiple logo variants with different styles"""
|
||||
|
||||
# Select appropriate styles for batch generation
|
||||
@@ -247,16 +427,20 @@ def generate_batch(prompt, brand_name, count, output_dir, use_pro=False, brand_c
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
results = []
|
||||
model_label = "Pro" if use_pro else "Flash"
|
||||
model_label = (
|
||||
f"Atlas Cloud ({atlas_model})"
|
||||
if provider == "atlas"
|
||||
else f"Nano Banana {'Pro' if use_pro else 'Flash'}"
|
||||
)
|
||||
ratio = aspect_ratio if aspect_ratio in ASPECT_RATIOS else DEFAULT_ASPECT_RATIO
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" BATCH LOGO GENERATION: {brand_name}")
|
||||
print(f" Model: Nano Banana {model_label}")
|
||||
print(f" Model: {model_label}")
|
||||
print(f" Aspect Ratio: {ratio}")
|
||||
print(f" Variants: {count}")
|
||||
print(f" Output: {output_dir}")
|
||||
print(f"{'='*60}\n")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
for i in range(min(count, len(batch_styles))):
|
||||
style_key, style_desc = batch_styles[i]
|
||||
@@ -267,10 +451,10 @@ def generate_batch(prompt, brand_name, count, output_dir, use_pro=False, brand_c
|
||||
enhanced_prompt = f"{brand_context}, {enhanced_prompt}"
|
||||
|
||||
# Generate filename
|
||||
filename = f"{brand_name.lower().replace(' ', '_')}_{style_key}_{i+1:02d}.png"
|
||||
filename = f"{brand_name.lower().replace(' ', '_')}_{style_key}_{i + 1:02d}.png"
|
||||
output_path = os.path.join(output_dir, filename)
|
||||
|
||||
print(f"[{i+1}/{count}] Generating {style_key} variant...")
|
||||
print(f"[{i + 1}/{count}] Generating {style_key} variant...")
|
||||
|
||||
result = generate_logo(
|
||||
prompt=enhanced_prompt,
|
||||
@@ -279,7 +463,9 @@ def generate_batch(prompt, brand_name, count, output_dir, use_pro=False, brand_c
|
||||
brand_name=brand_name,
|
||||
output_path=output_path,
|
||||
use_pro=use_pro,
|
||||
aspect_ratio=aspect_ratio
|
||||
aspect_ratio=aspect_ratio,
|
||||
provider=provider,
|
||||
atlas_model=atlas_model,
|
||||
)
|
||||
|
||||
if result:
|
||||
@@ -292,31 +478,70 @@ def generate_batch(prompt, brand_name, count, output_dir, use_pro=False, brand_c
|
||||
if i < count - 1:
|
||||
time.sleep(2)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" BATCH COMPLETE: {len(results)}/{count} logos generated")
|
||||
print(f"{'='*60}\n")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate logos using Gemini Nano Banana models")
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate logos using Gemini or Atlas Cloud"
|
||||
)
|
||||
parser.add_argument("--prompt", "-p", type=str, help="Logo description prompt")
|
||||
parser.add_argument("--brand", "-b", type=str, help="Brand name")
|
||||
parser.add_argument("--style", "-s", choices=list(STYLE_MODIFIERS.keys()), help="Logo style")
|
||||
parser.add_argument("--industry", "-i", choices=list(INDUSTRY_PROMPTS.keys()), help="Industry type")
|
||||
parser.add_argument(
|
||||
"--style", "-s", choices=list(STYLE_MODIFIERS.keys()), help="Logo style"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--industry", "-i", choices=list(INDUSTRY_PROMPTS.keys()), help="Industry type"
|
||||
)
|
||||
parser.add_argument("--output", "-o", type=str, help="Output file path")
|
||||
parser.add_argument("--output-dir", type=str, help="Output directory for batch generation")
|
||||
parser.add_argument("--batch", type=int, help="Number of logo variants to generate (batch mode)")
|
||||
parser.add_argument("--brand-context", type=str, help="Additional brand context for prompts")
|
||||
parser.add_argument("--pro", action="store_true", help="Use Nano Banana Pro (gemini-3-pro-image-preview) for professional quality")
|
||||
parser.add_argument("--aspect-ratio", "-r", choices=ASPECT_RATIOS, default=DEFAULT_ASPECT_RATIO,
|
||||
help=f"Image aspect ratio (default: {DEFAULT_ASPECT_RATIO} for logos)")
|
||||
parser.add_argument("--list-styles", action="store_true", help="List available styles")
|
||||
parser.add_argument("--list-industries", action="store_true", help="List available industries")
|
||||
parser.add_argument(
|
||||
"--output-dir", type=str, help="Output directory for batch generation"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch", type=int, help="Number of logo variants to generate (batch mode)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--brand-context", type=str, help="Additional brand context for prompts"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pro",
|
||||
action="store_true",
|
||||
help="Use Nano Banana Pro (gemini-3-pro-image-preview) for professional quality",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--provider",
|
||||
choices=["gemini", "atlas"],
|
||||
default="gemini",
|
||||
help="Image provider (default: gemini)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--atlas-model",
|
||||
default=ATLAS_MODEL,
|
||||
help=f"Atlas Cloud image model (default: {ATLAS_MODEL})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--aspect-ratio",
|
||||
"-r",
|
||||
choices=ASPECT_RATIOS,
|
||||
default=DEFAULT_ASPECT_RATIO,
|
||||
help=f"Image aspect ratio (default: {DEFAULT_ASPECT_RATIO} for logos)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--list-styles", action="store_true", help="List available styles"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--list-industries", action="store_true", help="List available industries"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.provider == "atlas" and args.pro:
|
||||
parser.error("--pro is only available with --provider gemini")
|
||||
|
||||
if args.list_styles:
|
||||
print("Available styles:")
|
||||
for style, desc in STYLE_MODIFIERS.items():
|
||||
@@ -336,7 +561,9 @@ def main():
|
||||
|
||||
# Batch mode
|
||||
if args.batch:
|
||||
output_dir = args.output_dir or f"./{args.brand.lower().replace(' ', '_')}_logos"
|
||||
output_dir = (
|
||||
args.output_dir or f"./{args.brand.lower().replace(' ', '_')}_logos"
|
||||
)
|
||||
generate_batch(
|
||||
prompt=prompt,
|
||||
brand_name=args.brand or "Logo",
|
||||
@@ -344,7 +571,9 @@ def main():
|
||||
output_dir=output_dir,
|
||||
use_pro=args.pro,
|
||||
brand_context=args.brand_context,
|
||||
aspect_ratio=args.aspect_ratio
|
||||
aspect_ratio=args.aspect_ratio,
|
||||
provider=args.provider,
|
||||
atlas_model=args.atlas_model,
|
||||
)
|
||||
else:
|
||||
generate_logo(
|
||||
@@ -354,7 +583,9 @@ def main():
|
||||
brand_name=args.brand,
|
||||
output_path=args.output,
|
||||
use_pro=args.pro,
|
||||
aspect_ratio=args.aspect_ratio
|
||||
aspect_ratio=args.aspect_ratio,
|
||||
provider=args.provider,
|
||||
atlas_model=args.atlas_model,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
import importlib.util
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import call, patch
|
||||
|
||||
MODULE_PATH = Path(__file__).parents[1] / "generate.py"
|
||||
SPEC = importlib.util.spec_from_file_location("logo_generate", MODULE_PATH)
|
||||
logo_generate = importlib.util.module_from_spec(SPEC)
|
||||
SPEC.loader.exec_module(logo_generate)
|
||||
|
||||
|
||||
class AtlasGenerationTests(unittest.TestCase):
|
||||
@patch.object(logo_generate, "_download_atlas_image")
|
||||
@patch.object(logo_generate.time, "sleep")
|
||||
@patch.object(logo_generate, "_json_request")
|
||||
def test_atlas_submits_once_and_polls_until_completed(
|
||||
self, json_request, sleep, download
|
||||
):
|
||||
json_request.side_effect = [
|
||||
{"code": 200, "data": {"id": "pred-123", "status": "created"}},
|
||||
{"code": 200, "data": {"id": "pred-123", "status": "processing"}},
|
||||
{
|
||||
"code": 200,
|
||||
"data": {
|
||||
"id": "pred-123",
|
||||
"status": "completed",
|
||||
"outputs": ["https://media.example.com/logo.png"],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
logo_generate._generate_with_atlas(
|
||||
"logo prompt", "logo.png", "1:1", "atlas-key", "atlas/model"
|
||||
)
|
||||
|
||||
self.assertEqual(json_request.call_count, 3)
|
||||
self.assertEqual(
|
||||
json_request.call_args_list[0],
|
||||
call(
|
||||
f"{logo_generate.ATLAS_API_BASE}/model/generateImage",
|
||||
"atlas-key",
|
||||
method="POST",
|
||||
payload={
|
||||
"model": "atlas/model",
|
||||
"prompt": "logo prompt",
|
||||
"aspect_ratio": "1:1",
|
||||
},
|
||||
),
|
||||
)
|
||||
self.assertEqual(
|
||||
json_request.call_args_list[1:],
|
||||
[
|
||||
call(
|
||||
f"{logo_generate.ATLAS_API_BASE}/model/prediction/pred-123",
|
||||
"atlas-key",
|
||||
),
|
||||
call(
|
||||
f"{logo_generate.ATLAS_API_BASE}/model/prediction/pred-123",
|
||||
"atlas-key",
|
||||
),
|
||||
],
|
||||
)
|
||||
self.assertEqual(sleep.call_count, 2)
|
||||
download.assert_called_once_with(
|
||||
"https://media.example.com/logo.png", "logo.png"
|
||||
)
|
||||
|
||||
@patch.object(logo_generate, "_json_request")
|
||||
def test_atlas_does_not_retry_generation_post(self, json_request):
|
||||
json_request.side_effect = RuntimeError("network error")
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "network error"):
|
||||
logo_generate._generate_with_atlas(
|
||||
"logo prompt", "logo.png", "1:1", "atlas-key", "atlas/model"
|
||||
)
|
||||
|
||||
json_request.assert_called_once()
|
||||
|
||||
@patch.object(logo_generate, "_validate_public_https_url")
|
||||
@patch.object(logo_generate, "build_opener")
|
||||
def test_media_download_never_forwards_api_key(self, build_opener, validate):
|
||||
class Headers:
|
||||
@staticmethod
|
||||
def get_content_type():
|
||||
return "image/png"
|
||||
|
||||
class Response:
|
||||
headers = Headers()
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def read():
|
||||
return b"png-bytes"
|
||||
|
||||
build_opener.return_value.open.return_value = Response()
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
output = Path(temp_dir) / "logo.png"
|
||||
logo_generate._download_atlas_image(
|
||||
"https://media.example.com/logo.png", output
|
||||
)
|
||||
self.assertEqual(output.read_bytes(), b"png-bytes")
|
||||
|
||||
request = build_opener.return_value.open.call_args.args[0]
|
||||
headers = {key.lower(): value for key, value in request.header_items()}
|
||||
self.assertNotIn("authorization", headers)
|
||||
self.assertEqual(headers["accept"], "image/*")
|
||||
self.assertEqual(headers["user-agent"], logo_generate.HTTP_USER_AGENT)
|
||||
validate.assert_called_once_with("https://media.example.com/logo.png")
|
||||
|
||||
def test_atlas_requires_api_key(self):
|
||||
with self.assertRaisesRegex(RuntimeError, "ATLASCLOUD_API_KEY not set"):
|
||||
logo_generate._generate_with_atlas(
|
||||
"logo prompt", "logo.png", "1:1", None, "atlas/model"
|
||||
)
|
||||
|
||||
def test_media_url_rejects_private_addresses(self):
|
||||
with self.assertRaisesRegex(ValueError, "non-public address"):
|
||||
logo_generate._validate_public_https_url("https://127.0.0.1/logo.png")
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "local hostname"):
|
||||
logo_generate._validate_public_https_url("https://assets.local/logo.png")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user