Buckets:

sharktide
/

tools

Files

xet

sharktide/tools / createpics.py

sharktide

3 days ago

download

raw

19.8 kB

#!/usr/bin/env python3
"""
Generates unique branded cinematic images for models from two sources:
  - Cloud models:  https://api.inferenceport.ai/v1/models   (OpenAI-style {"data":[{"id": ...}]})
  - Local models:  https://api.inferenceport.ai/models       (Ollama-style [{"name": ...}, ...])
using the Pollinations image API (model=zimage).
Naming convention:
  Cloud model "foo/bar":
    {OUTPUT_DIR}/foo_bar.jpg
    {OUTPUT_DIR}/compressed/foo_bar.webp
  Local model "foo/bar":
    {OUTPUT_DIR}/foo_bar.local.jpg
    {OUTPUT_DIR}/compressed/foo_bar.local.webp
Each model gets a stable, art-directed concept:
a bold branded scene (fog, neon, metallic text, surreal landscapes,
cosmic backdrops, etc.) with the model's name rendered as giant stylized
text integrated into the scene. The prompt generator is seeded off the
model's own name, so the same model always gets the same scene, and
different models get different combinations of setting / lighting /
text-material / camera style / palette.
USAGE
-----
Scan both endpoints and fill in anything missing (default command):
    python tools/createpics.py
    python tools/createpics.py scan
Create an image for a single new model:
    python tools/createpics.py create -n "some-model-id" -t cloud
    python tools/createpics.py create -n "llama3:8b" -t local
Regenerate (overwrite) an existing model's image, e.g. if the text
came out warped:
    python tools/createpics.py regenerate -n "some-model-id" -t cloud
"""

import os
import re
import time
import random
import hashlib
import argparse
import urllib.parse
from pathlib import Path
from huggingface_hub import sync_bucket
import requests
from rich.progress import (
    Progress,
    BarColumn,
    TextColumn,
    TimeElapsedColumn,
    MofNCompleteColumn,
)
from rich.console import Console

REPO_ROOT = Path(__file__).resolve().parent.parent

# --------------------------------------------------------------------------
# CONFIG
# --------------------------------------------------------------------------
DEFAULT_API_KEY = "sk_w1UrhHef25QNpFGEFnOGanL4WIaJLbVI"
MODELS_ENDPOINT = "https://api.inferenceport.ai/models"
V1_MODELS_ENDPOINT = "https://us-east-2.api.inferenceport.ai/v1/models"
POLLINATIONS_URL_TEMPLATE = "https://gen.pollinations.ai/image/{prompt}?model=zimage&key={key}&seed={seed}"
REQUEST_TIMEOUT = 60
RETRY_COUNT = 2
RETRY_DELAY = 3
WEBP_QUALITY = 82

console = Console()

# --------------------------------------------------------------------------
# ART DIRECTION LIBRARY
# A cohesive concept is selected before the palette and framing.
# --------------------------------------------------------------------------

ART_DIRECTIONS = [
    ("architectural editorial", "an impossible brutalist museum suspended over a quiet sea, with the name carved across the main cantilever", "rigorous geometry, soft film grain"),
    ("retro space poster", "a hand-painted seventies spaceport beneath an enormous setting sun, with the name set as the poster masthead", "vintage airbrush and screenprint"),
    ("surreal botanical", "a glasshouse of gigantic translucent orchids and mineral specimens, with the name grown as botanical lettering", "fine-art botanical photography"),
    ("fashion still life", "folded silk, polished stone and floating ribbons forming the name", "high-fashion magazine photography"),
    ("kinetic installation", "an immense mobile sculpture of interlocking rings and suspended letterforms spelling the name", "contemporary museum installation"),
    ("ocean archaeology", "a submersible discovering an ancient artifact engraved with the name on the ocean floor", "underwater documentary"),
    ("clay miniature", "a miniature ceramic world of whimsical roads and towers with the name assembled from handmade clay", "stop-motion set photography"),
    ("typographic city", "a dream metropolis with elevated trains threading between letter-shaped buildings spelling the name", "isometric architectural illustration"),
    ("celestial atlas", "a beautifully engraved astronomical chart of imaginary moons and orbital paths, bearing the name as its main heading", "antique scientific print"),
    ("desert land art", "colossal trenches and sculpted dunes spelling the name as an aerial earthwork", "fine-art aerial photography"),
    ("optics laboratory", "a precision laboratory of floating glass prisms, with the name etched on its main lens", "luxury product photography"),
    ("storybook", "a village built on the back of a wandering mountain, with the name painted across a grand sign", "lavish storybook illustration"),
    ("fossil museum", "a mysterious geometric fossil displayed on a plinth engraved with the name", "natural-history exhibition photography"),
    ("minimal product", "an unusual machined object floating above a matte pedestal beneath a crisp sign bearing the name", "minimal industrial design campaign"),
    ("paper theatre", "a layered paper-cut mountain railway with the name worked into its proscenium", "handcrafted paper sculpture"),
    ("graphic novel", "a dramatic rooftop scene with sweeping lines and the name as bold in-world comic lettering", "inked graphic novel splash page"),
    ("impossible library", "an Escher-like archive of floating books, with the name embossed on a giant volume", "magical-realist film still"),
    ("arcade nostalgia", "an original arcade machine whose illuminated marquee spells the name", "retro commercial photography"),
    ("transit diagram", "an imaginary metro network and miniature stations with the name as the system's graphic identity", "Swiss information design"),
    ("porcelain world", "a landscape made of glazed porcelain hills and bridges, with raised ceramic lettering spelling the name", "studio ceramics photography"),
    ("festival installation", "a vast floating fabric sculpture over a public square with the name integrated into suspended signage", "contemporary art photography"),
    ("mineral macro", "impossible gemstone formations with the name engraved into a luminous crystal face", "extreme macro photography"),
    ("cinematic railway", "a futuristic train crossing a bridge above thick clouds, the name painted along its carriages", "cinematic wide establishing shot"),
    ("noir storefront", "a rain-soaked shop of mysterious inventions, with the name as its elegant exterior signage", "moody narrative film still"),
]

PALETTES = [
    "burnt orange, ink blue and cream", "copper, sea glass and charcoal",
    "cherry red, butter yellow and ultramarine", "lilac, slate and silver",
    "forest green, parchment and terracotta", "cobalt, coral and black",
    "champagne, plum and white", "sunflower yellow, gray and cyan",
    "burgundy, pale pink and midnight blue", "ivory and graphite with vermilion",
    "mint, dark chocolate and apricot", "lapis blue, bronze and sand",
]
FRAMINGS = [
    "wide panorama with strong visual hierarchy", "balanced centered composition",
    "off-center framing with generous negative space", "dramatic low-angle view",
    "striking overhead composition", "three-quarter view with layered depth",
]

def build_prompt(display_name: str) -> str:
    """Give each model a stable, distinctive art direction with legible branding."""
    rng = random.Random(int.from_bytes(hashlib.sha256(display_name.encode("utf-8")).digest(), "big"))
    genre, scene, style = rng.choice(ART_DIRECTIONS)
    palette, framing = rng.choice(PALETTES), rng.choice(FRAMINGS)
    name = display_name.upper()
    return (
        f"Original art-directed cover image for an AI model named \"{name}\". "
        f"Direction: {genre}. Depict {scene.replace('the name', repr(name))}. "
        f"Art style: {style}; {framing}. Deliberate palette: {palette}. "
        f"Render the exact text \"{name}\" prominently, clearly and correctly spelled; "
        "integrate it physically into the idea rather than applying a generic overlay. "
        "Deliver an instantly memorable central concept with coherent material, "
        "lighting and composition. No other words, logos or watermarks. "
        "Avoid interchangeable neon landscapes and giant chrome titles."
    )


def safe_filename(name: str) -> str:
    """Turn a model name (which may contain slashes, dots, colons, etc.) into a safe filename."""
    return re.sub(r"[^A-Za-z0-9._-]", "_", name)


def paths_for(model_id: str, model_type: str, output_dir: str):
    """Return (jpg_path, webp_path) for a given model id/type."""
    safe = safe_filename(model_id)
    if model_type == "local":
        jpg_name = f"{safe}.local.jpg"
        webp_name = f"{safe}.local.webp"
    else:
        jpg_name = f"{safe}.jpg"
        webp_name = f"{safe}.webp"
    jpg_path = os.path.join(output_dir, jpg_name)
    webp_path = os.path.join(output_dir, "compressed", webp_name)
    return jpg_path, webp_path


# --------------------------------------------------------------------------
# FETCHING MODEL LISTS
# --------------------------------------------------------------------------

def fetch_local_models() -> list[str]:
    """Fetch the /models (Ollama-style) list and return model names."""
    response = requests.get(MODELS_ENDPOINT, timeout=REQUEST_TIMEOUT)
    response.raise_for_status()
    data = response.json()

    names = []
    for entry in data:
        name = entry.get("name")
        if name:
            names.append(name)
    return names


def fetch_cloud_models() -> list[str]:
    """Fetch the /v1/models (OpenAI-style) list and return model ids."""
    response = requests.get(V1_MODELS_ENDPOINT, timeout=REQUEST_TIMEOUT)
    response.raise_for_status()
    data = response.json()

    items = data.get("data", data) if isinstance(data, dict) else data

    ids = []
    for entry in items:
        model_id = entry.get("id") or entry.get("name")
        if model_id:
            ids.append(model_id)
    return ids


# --------------------------------------------------------------------------
# IMAGE GENERATION / CONVERSION
# --------------------------------------------------------------------------

def fetch_and_save_image(model_id: str, prompt: str, jpg_path: str, api_key: str, seed: int = 0) -> bool:
    """Fetch one image from Pollinations and save it to jpg_path. Returns success bool."""
    encoded_prompt = urllib.parse.quote(prompt, safe="")
    url = POLLINATIONS_URL_TEMPLATE.format(prompt=encoded_prompt, key=api_key, seed=seed)

    os.makedirs(os.path.dirname(jpg_path) or ".", exist_ok=True)

    last_error = None
    for attempt in range(1, RETRY_COUNT + 2):  # initial try + retries
        try:
            response = requests.get(url, timeout=REQUEST_TIMEOUT)
            response.raise_for_status()

            content_type = response.headers.get("Content-Type", "")
            if "image" not in content_type and len(response.content) < 1000:
                raise ValueError(
                    f"Unexpected response (content-type={content_type!r}, "
                    f"{len(response.content)} bytes): {response.content[:200]!r}"
                )

            with open(jpg_path, "wb") as f:
                f.write(response.content)
            return True

        except Exception as e:  # noqa: BLE001 - catch & report everything
            last_error = e
            if attempt <= RETRY_COUNT:
                time.sleep(RETRY_DELAY)

    console.print(f"[bold red]✗ Failed to generate:[/bold red] {model_id} — {last_error}")
    return False


def convert_to_webp(jpg_path: str, webp_path: str, quality: int = WEBP_QUALITY) -> bool:
    """Convert an existing jpg to webp. Returns success bool."""
    try:
        from PIL import Image
    except ImportError:
        console.print(
            "[bold red]Pillow is required for webp conversion.[/bold red] "
            "Install it with: pip install pillow --break-system-packages"
        )
        return False

    os.makedirs(os.path.dirname(webp_path) or ".", exist_ok=True)
    try:
        with Image.open(jpg_path) as im:
            im = im.convert("RGB")
            im.save(webp_path, "WEBP", quality=quality)
        return True
    except Exception as e:  # noqa: BLE001
        console.print(f"[bold red]✗ WebP conversion failed:[/bold red] {jpg_path} — {e}")
        return False


def process_model(model_id: str, model_type: str, output_dir: str, api_key: str, overwrite: bool = False, seed: int = 0) -> bool:
    """Generate (if needed) and convert (if needed) images for a single model."""
    jpg_path, webp_path = paths_for(model_id, model_type, output_dir)
    prompt = build_prompt(model_id)

    need_jpg = overwrite or not os.path.exists(jpg_path)
    if need_jpg:
        ok = fetch_and_save_image(model_id, prompt, jpg_path, api_key, seed=seed)
        if not ok:
            return False

    need_webp = overwrite or not os.path.exists(webp_path)
    if need_webp:
        convert_to_webp(jpg_path, webp_path)

    return True


# --------------------------------------------------------------------------
# COMMANDS
# --------------------------------------------------------------------------

def cmd_scan(output_dir: str, api_key: str):
    console.print("[bold cyan]Fetching model lists...[/bold cyan]")

    try:
        cloud_models = fetch_cloud_models()
    except Exception as e:
        console.print(f"[bold red]Failed to fetch /v1/models:[/bold red] {e}")
        cloud_models = []

    try:
        local_models = fetch_local_models()
    except Exception as e:
        console.print(f"[bold red]Failed to fetch /models:[/bold red] {e}")
        local_models = []

    console.print(
        f"[bold cyan]Found {len(cloud_models)} cloud models and "
        f"{len(local_models)} local models.[/bold cyan]"
    )

    all_jobs = [(m, "cloud") for m in cloud_models] + [(m, "local") for m in local_models]

    # Only queue work for models missing a jpg and/or webp.
    pending = []
    for model_id, model_type in all_jobs:
        jpg_path, webp_path = paths_for(model_id, model_type, output_dir)
        if not os.path.exists(jpg_path) or not os.path.exists(webp_path):
            pending.append((model_id, model_type))

    console.print(f"[bold cyan]{len(pending)} image(s) need generating/converting.[/bold cyan]\n")

    if not pending:
        console.print("[bold green]Everything is already up to date.[/bold green]")
        return

    with Progress(
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        MofNCompleteColumn(),
        TimeElapsedColumn(),
        console=console,
    ) as progress:
        task = progress.add_task("Generating images", total=len(pending))

        for model_id, model_type in pending:
            progress.update(task, description=f"[green]{model_id} ({model_type})[/green]")
            process_model(model_id, model_type, output_dir, api_key)
            progress.advance(task)

    console.print(f"\n[bold green]Done.[/bold green] Images saved under [bold]{output_dir}[/bold]")


def cmd_create(name: str, model_type: str, output_dir: str, api_key: str):
    jpg_path, webp_path = paths_for(name, model_type, output_dir)

    if os.path.exists(jpg_path) and os.path.exists(webp_path):
        console.print(
            f"[bold yellow]Already exists:[/bold yellow] {jpg_path}\n"
            f"Use the 'regenerate' command if you want to overwrite it."
        )
        return

    console.print(f"[bold cyan]Creating image for[/bold cyan] {name} [dim]({model_type})[/dim]")
    ok = process_model(name, model_type, output_dir, api_key, overwrite=False)
    if ok:
        console.print(f"[bold green]Saved:[/bold green] {jpg_path}\n[bold green]Saved:[/bold green] {webp_path}")
    else:
        console.print("[bold red]Failed to create image.[/bold red]")


def cmd_regenerate(name: str, model_type: str, output_dir: str, api_key: str):
    jpg_path, webp_path = paths_for(name, model_type, output_dir)

    console.print(f"[bold cyan]Regenerating image for[/bold cyan] {name} [dim]({model_type})[/dim]")
    ok = process_model(name, model_type, output_dir, api_key, overwrite=True, seed=random.randint(1, 200))
    if ok:
        console.print(f"[bold green]Overwrote:[/bold green] {jpg_path}\n[bold green]Overwrote:[/bold green] {webp_path}")
    else:
        console.print("[bold red]Failed to regenerate image.[/bold red]")


# --------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------

def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Generate branded images for inference models (cloud + local)."
    )
    parser.add_argument("--api-key", help="Your Pollinations API key.")
    parser.add_argument("--output-dir", default=str(REPO_ROOT / "modelpics"), help="Directory to save generated images.")

    subparsers = parser.add_subparsers(dest="command")

    subparsers.add_parser(
        "scan",
        help="Scan both endpoints and generate any missing images/webp conversions (default).",
    )

    create_p = subparsers.add_parser("create", help="Create an image for a single new model.")
    create_p.add_argument("-n", "--name", required=True, help="The model ID/name.")
    create_p.add_argument("-t", "--type", required=True, choices=["cloud", "local"], help="Model source type.")

    regen_p = subparsers.add_parser("regenerate", help="Regenerate (overwrite) an existing model's image.")
    regen_p.add_argument("-n", "--name", required=True, help="The model ID/name.")
    regen_p.add_argument("-t", "--type", required=True, choices=["cloud", "local"], help="Model source type.")

    subparsers.add_parser("sync", help="Sync the local output directory with the HF bucket.")
    return parser

def get_md5(file_path):
    """Generate a quick hash of the file content."""
    hasher = hashlib.md5()
    with open(file_path, 'rb') as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hasher.update(chunk)
    return hasher.hexdigest()
def clean_local_renames_before_sync(local_dir):
    """Finds duplicate files by content and removes the older-named version."""
    seen_hashes = {}
    for root, _, files in os.walk(local_dir):
        for file in files:
            path = os.path.join(root, file)
            f_hash = get_md5(path)
            
            if f_hash in seen_hashes:
                # We found a rename/duplicate! Delete the ghost/old copy
                print(f"⚠️ Detected rename conflict. Removing local ghost file: {path}")
                os.remove(path)
            else:
                seen_hashes[f_hash] = path

def sync(output_dir):
    clean_local_renames_before_sync(output_dir)

    sync_bucket("hf://buckets/inferenceport-ai/model-info", output_dir, delete=False)
    
    sync_bucket(output_dir, "hf://buckets/inferenceport-ai/model-info", delete=False)

def main():
    parser = build_arg_parser()
    args = parser.parse_args()

    api_key = args.api_key or DEFAULT_API_KEY
    output_dir = args.output_dir

    if api_key in ("YOUR_POLLINATIONS_API_KEY_HERE", "", None):
        console.print(
            "[bold yellow]Warning:[/bold yellow] API_KEY is still a placeholder. "
            "Pass --api-key or set DEFAULT_API_KEY before running.\n"
        )

    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(os.path.join(output_dir, "compressed"), exist_ok=True)

    command = args.command or "scan"

    if command == "scan":
        cmd_scan(output_dir, api_key)
    elif command == "create":
        cmd_create(args.name, args.type, output_dir, api_key)
    elif command == "regenerate":
        cmd_regenerate(args.name, args.type, output_dir, api_key)
    elif command == "sync":
        sync(output_dir)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 19.8 kB
Xet hash:: 9968a2bfd3356b15dcb16ef7449c350948d90f6552c6a834c68385ed29519996

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	#!/usr/bin/env python3
	"""
	Generates unique branded cinematic images for models from two sources:

	- Cloud models: https://api.inferenceport.ai/v1/models (OpenAI-style {"data":[{"id": ...}]})
	- Local models: https://api.inferenceport.ai/models (Ollama-style [{"name": ...}, ...])

	using the Pollinations image API (model=zimage).

	Naming convention:
	Cloud model "foo/bar":
	{OUTPUT_DIR}/foo_bar.jpg
	{OUTPUT_DIR}/compressed/foo_bar.webp
	Local model "foo/bar":
	{OUTPUT_DIR}/foo_bar.local.jpg
	{OUTPUT_DIR}/compressed/foo_bar.local.webp

	Each model gets a stable, art-directed concept:
	a bold branded scene (fog, neon, metallic text, surreal landscapes,
	cosmic backdrops, etc.) with the model's name rendered as giant stylized
	text integrated into the scene. The prompt generator is seeded off the
	model's own name, so the same model always gets the same scene, and
	different models get different combinations of setting / lighting /
	text-material / camera style / palette.

	USAGE
	-----
	Scan both endpoints and fill in anything missing (default command):
	python tools/createpics.py
	python tools/createpics.py scan

	Create an image for a single new model:
	python tools/createpics.py create -n "some-model-id" -t cloud
	python tools/createpics.py create -n "llama3:8b" -t local

	Regenerate (overwrite) an existing model's image, e.g. if the text
	came out warped:
	python tools/createpics.py regenerate -n "some-model-id" -t cloud
	"""

	import os
	import re
	import time
	import random
	import hashlib
	import argparse
	import urllib.parse
	from pathlib import Path
	from huggingface_hub import sync_bucket
	import requests
	from rich.progress import (
	Progress,
	BarColumn,
	TextColumn,
	TimeElapsedColumn,
	MofNCompleteColumn,
	)
	from rich.console import Console

	REPO_ROOT = Path(__file__).resolve().parent.parent

	# --------------------------------------------------------------------------
	# CONFIG
	# --------------------------------------------------------------------------
	DEFAULT_API_KEY = "sk_w1UrhHef25QNpFGEFnOGanL4WIaJLbVI"
	MODELS_ENDPOINT = "https://api.inferenceport.ai/models"
	V1_MODELS_ENDPOINT = "https://us-east-2.api.inferenceport.ai/v1/models"
	POLLINATIONS_URL_TEMPLATE = "https://gen.pollinations.ai/image/{prompt}?model=zimage&key={key}&seed={seed}"
	REQUEST_TIMEOUT = 60
	RETRY_COUNT = 2
	RETRY_DELAY = 3
	WEBP_QUALITY = 82

	console = Console()

	# --------------------------------------------------------------------------
	# ART DIRECTION LIBRARY
	# A cohesive concept is selected before the palette and framing.
	# --------------------------------------------------------------------------

	ART_DIRECTIONS = [
	("architectural editorial", "an impossible brutalist museum suspended over a quiet sea, with the name carved across the main cantilever", "rigorous geometry, soft film grain"),
	("retro space poster", "a hand-painted seventies spaceport beneath an enormous setting sun, with the name set as the poster masthead", "vintage airbrush and screenprint"),
	("surreal botanical", "a glasshouse of gigantic translucent orchids and mineral specimens, with the name grown as botanical lettering", "fine-art botanical photography"),
	("fashion still life", "folded silk, polished stone and floating ribbons forming the name", "high-fashion magazine photography"),
	("kinetic installation", "an immense mobile sculpture of interlocking rings and suspended letterforms spelling the name", "contemporary museum installation"),
	("ocean archaeology", "a submersible discovering an ancient artifact engraved with the name on the ocean floor", "underwater documentary"),
	("clay miniature", "a miniature ceramic world of whimsical roads and towers with the name assembled from handmade clay", "stop-motion set photography"),
	("typographic city", "a dream metropolis with elevated trains threading between letter-shaped buildings spelling the name", "isometric architectural illustration"),
	("celestial atlas", "a beautifully engraved astronomical chart of imaginary moons and orbital paths, bearing the name as its main heading", "antique scientific print"),
	("desert land art", "colossal trenches and sculpted dunes spelling the name as an aerial earthwork", "fine-art aerial photography"),
	("optics laboratory", "a precision laboratory of floating glass prisms, with the name etched on its main lens", "luxury product photography"),
	("storybook", "a village built on the back of a wandering mountain, with the name painted across a grand sign", "lavish storybook illustration"),
	("fossil museum", "a mysterious geometric fossil displayed on a plinth engraved with the name", "natural-history exhibition photography"),
	("minimal product", "an unusual machined object floating above a matte pedestal beneath a crisp sign bearing the name", "minimal industrial design campaign"),
	("paper theatre", "a layered paper-cut mountain railway with the name worked into its proscenium", "handcrafted paper sculpture"),
	("graphic novel", "a dramatic rooftop scene with sweeping lines and the name as bold in-world comic lettering", "inked graphic novel splash page"),
	("impossible library", "an Escher-like archive of floating books, with the name embossed on a giant volume", "magical-realist film still"),
	("arcade nostalgia", "an original arcade machine whose illuminated marquee spells the name", "retro commercial photography"),
	("transit diagram", "an imaginary metro network and miniature stations with the name as the system's graphic identity", "Swiss information design"),
	("porcelain world", "a landscape made of glazed porcelain hills and bridges, with raised ceramic lettering spelling the name", "studio ceramics photography"),
	("festival installation", "a vast floating fabric sculpture over a public square with the name integrated into suspended signage", "contemporary art photography"),
	("mineral macro", "impossible gemstone formations with the name engraved into a luminous crystal face", "extreme macro photography"),
	("cinematic railway", "a futuristic train crossing a bridge above thick clouds, the name painted along its carriages", "cinematic wide establishing shot"),
	("noir storefront", "a rain-soaked shop of mysterious inventions, with the name as its elegant exterior signage", "moody narrative film still"),
	]

	PALETTES = [
	"burnt orange, ink blue and cream", "copper, sea glass and charcoal",
	"cherry red, butter yellow and ultramarine", "lilac, slate and silver",
	"forest green, parchment and terracotta", "cobalt, coral and black",
	"champagne, plum and white", "sunflower yellow, gray and cyan",
	"burgundy, pale pink and midnight blue", "ivory and graphite with vermilion",
	"mint, dark chocolate and apricot", "lapis blue, bronze and sand",
	]
	FRAMINGS = [
	"wide panorama with strong visual hierarchy", "balanced centered composition",
	"off-center framing with generous negative space", "dramatic low-angle view",
	"striking overhead composition", "three-quarter view with layered depth",
	]

	def build_prompt(display_name: str) -> str:
	"""Give each model a stable, distinctive art direction with legible branding."""
	rng = random.Random(int.from_bytes(hashlib.sha256(display_name.encode("utf-8")).digest(), "big"))
	genre, scene, style = rng.choice(ART_DIRECTIONS)
	palette, framing = rng.choice(PALETTES), rng.choice(FRAMINGS)
	name = display_name.upper()
	return (
	f"Original art-directed cover image for an AI model named \"{name}\". "
	f"Direction: {genre}. Depict {scene.replace('the name', repr(name))}. "
	f"Art style: {style}; {framing}. Deliberate palette: {palette}. "
	f"Render the exact text \"{name}\" prominently, clearly and correctly spelled; "
	"integrate it physically into the idea rather than applying a generic overlay. "
	"Deliver an instantly memorable central concept with coherent material, "
	"lighting and composition. No other words, logos or watermarks. "
	"Avoid interchangeable neon landscapes and giant chrome titles."
	)


	def safe_filename(name: str) -> str:
	"""Turn a model name (which may contain slashes, dots, colons, etc.) into a safe filename."""
	return re.sub(r"[^A-Za-z0-9._-]", "_", name)


	def paths_for(model_id: str, model_type: str, output_dir: str):
	"""Return (jpg_path, webp_path) for a given model id/type."""
	safe = safe_filename(model_id)
	if model_type == "local":
	jpg_name = f"{safe}.local.jpg"
	webp_name = f"{safe}.local.webp"
	else:
	jpg_name = f"{safe}.jpg"
	webp_name = f"{safe}.webp"
	jpg_path = os.path.join(output_dir, jpg_name)
	webp_path = os.path.join(output_dir, "compressed", webp_name)
	return jpg_path, webp_path


	# --------------------------------------------------------------------------
	# FETCHING MODEL LISTS
	# --------------------------------------------------------------------------

	def fetch_local_models() -> list[str]:
	"""Fetch the /models (Ollama-style) list and return model names."""
	response = requests.get(MODELS_ENDPOINT, timeout=REQUEST_TIMEOUT)
	response.raise_for_status()
	data = response.json()

	names = []
	for entry in data:
	name = entry.get("name")
	if name:
	names.append(name)
	return names


	def fetch_cloud_models() -> list[str]:
	"""Fetch the /v1/models (OpenAI-style) list and return model ids."""
	response = requests.get(V1_MODELS_ENDPOINT, timeout=REQUEST_TIMEOUT)
	response.raise_for_status()
	data = response.json()

	items = data.get("data", data) if isinstance(data, dict) else data

	ids = []
	for entry in items:
	model_id = entry.get("id") or entry.get("name")
	if model_id:
	ids.append(model_id)
	return ids


	# --------------------------------------------------------------------------
	# IMAGE GENERATION / CONVERSION
	# --------------------------------------------------------------------------

	def fetch_and_save_image(model_id: str, prompt: str, jpg_path: str, api_key: str, seed: int = 0) -> bool:
	"""Fetch one image from Pollinations and save it to jpg_path. Returns success bool."""
	encoded_prompt = urllib.parse.quote(prompt, safe="")
	url = POLLINATIONS_URL_TEMPLATE.format(prompt=encoded_prompt, key=api_key, seed=seed)

	os.makedirs(os.path.dirname(jpg_path) or ".", exist_ok=True)

	last_error = None
	for attempt in range(1, RETRY_COUNT + 2): # initial try + retries
	try:
	response = requests.get(url, timeout=REQUEST_TIMEOUT)
	response.raise_for_status()

	content_type = response.headers.get("Content-Type", "")
	if "image" not in content_type and len(response.content) < 1000:
	raise ValueError(
	f"Unexpected response (content-type={content_type!r}, "
	f"{len(response.content)} bytes): {response.content[:200]!r}"
	)

	with open(jpg_path, "wb") as f:
	f.write(response.content)
	return True

	except Exception as e: # noqa: BLE001 - catch & report everything
	last_error = e
	if attempt <= RETRY_COUNT:
	time.sleep(RETRY_DELAY)

	console.print(f"[bold red]✗ Failed to generate:[/bold red] {model_id} — {last_error}")
	return False


	def convert_to_webp(jpg_path: str, webp_path: str, quality: int = WEBP_QUALITY) -> bool:
	"""Convert an existing jpg to webp. Returns success bool."""
	try:
	from PIL import Image
	except ImportError:
	console.print(
	"[bold red]Pillow is required for webp conversion.[/bold red] "
	"Install it with: pip install pillow --break-system-packages"
	)
	return False

	os.makedirs(os.path.dirname(webp_path) or ".", exist_ok=True)
	try:
	with Image.open(jpg_path) as im:
	im = im.convert("RGB")
	im.save(webp_path, "WEBP", quality=quality)
	return True
	except Exception as e: # noqa: BLE001
	console.print(f"[bold red]✗ WebP conversion failed:[/bold red] {jpg_path} — {e}")
	return False


	def process_model(model_id: str, model_type: str, output_dir: str, api_key: str, overwrite: bool = False, seed: int = 0) -> bool:
	"""Generate (if needed) and convert (if needed) images for a single model."""
	jpg_path, webp_path = paths_for(model_id, model_type, output_dir)
	prompt = build_prompt(model_id)

	need_jpg = overwrite or not os.path.exists(jpg_path)
	if need_jpg:
	ok = fetch_and_save_image(model_id, prompt, jpg_path, api_key, seed=seed)
	if not ok:
	return False

	need_webp = overwrite or not os.path.exists(webp_path)
	if need_webp:
	convert_to_webp(jpg_path, webp_path)

	return True


	# --------------------------------------------------------------------------
	# COMMANDS
	# --------------------------------------------------------------------------

	def cmd_scan(output_dir: str, api_key: str):
	console.print("[bold cyan]Fetching model lists...[/bold cyan]")

	try:
	cloud_models = fetch_cloud_models()
	except Exception as e:
	console.print(f"[bold red]Failed to fetch /v1/models:[/bold red] {e}")
	cloud_models = []

	try:
	local_models = fetch_local_models()
	except Exception as e:
	console.print(f"[bold red]Failed to fetch /models:[/bold red] {e}")
	local_models = []

	console.print(
	f"[bold cyan]Found {len(cloud_models)} cloud models and "
	f"{len(local_models)} local models.[/bold cyan]"
	)

	all_jobs = [(m, "cloud") for m in cloud_models] + [(m, "local") for m in local_models]

	# Only queue work for models missing a jpg and/or webp.
	pending = []
	for model_id, model_type in all_jobs:
	jpg_path, webp_path = paths_for(model_id, model_type, output_dir)
	if not os.path.exists(jpg_path) or not os.path.exists(webp_path):
	pending.append((model_id, model_type))

	console.print(f"[bold cyan]{len(pending)} image(s) need generating/converting.[/bold cyan]\n")

	if not pending:
	console.print("[bold green]Everything is already up to date.[/bold green]")
	return

	with Progress(
	TextColumn("[progress.description]{task.description}"),
	BarColumn(),
	MofNCompleteColumn(),
	TimeElapsedColumn(),
	console=console,
	) as progress:
	task = progress.add_task("Generating images", total=len(pending))

	for model_id, model_type in pending:
	progress.update(task, description=f"[green]{model_id} ({model_type})[/green]")
	process_model(model_id, model_type, output_dir, api_key)
	progress.advance(task)

	console.print(f"\n[bold green]Done.[/bold green] Images saved under [bold]{output_dir}[/bold]")


	def cmd_create(name: str, model_type: str, output_dir: str, api_key: str):
	jpg_path, webp_path = paths_for(name, model_type, output_dir)

	if os.path.exists(jpg_path) and os.path.exists(webp_path):
	console.print(
	f"[bold yellow]Already exists:[/bold yellow] {jpg_path}\n"
	f"Use the 'regenerate' command if you want to overwrite it."
	)
	return

	console.print(f"[bold cyan]Creating image for[/bold cyan] {name} [dim]({model_type})[/dim]")
	ok = process_model(name, model_type, output_dir, api_key, overwrite=False)
	if ok:
	console.print(f"[bold green]Saved:[/bold green] {jpg_path}\n[bold green]Saved:[/bold green] {webp_path}")
	else:
	console.print("[bold red]Failed to create image.[/bold red]")


	def cmd_regenerate(name: str, model_type: str, output_dir: str, api_key: str):
	jpg_path, webp_path = paths_for(name, model_type, output_dir)

	console.print(f"[bold cyan]Regenerating image for[/bold cyan] {name} [dim]({model_type})[/dim]")
	ok = process_model(name, model_type, output_dir, api_key, overwrite=True, seed=random.randint(1, 200))
	if ok:
	console.print(f"[bold green]Overwrote:[/bold green] {jpg_path}\n[bold green]Overwrote:[/bold green] {webp_path}")
	else:
	console.print("[bold red]Failed to regenerate image.[/bold red]")


	# --------------------------------------------------------------------------
	# CLI
	# --------------------------------------------------------------------------

	def build_arg_parser() -> argparse.ArgumentParser:
	parser = argparse.ArgumentParser(
	description="Generate branded images for inference models (cloud + local)."
	)
	parser.add_argument("--api-key", help="Your Pollinations API key.")
	parser.add_argument("--output-dir", default=str(REPO_ROOT / "modelpics"), help="Directory to save generated images.")

	subparsers = parser.add_subparsers(dest="command")

	subparsers.add_parser(
	"scan",
	help="Scan both endpoints and generate any missing images/webp conversions (default).",
	)

	create_p = subparsers.add_parser("create", help="Create an image for a single new model.")
	create_p.add_argument("-n", "--name", required=True, help="The model ID/name.")
	create_p.add_argument("-t", "--type", required=True, choices=["cloud", "local"], help="Model source type.")

	regen_p = subparsers.add_parser("regenerate", help="Regenerate (overwrite) an existing model's image.")
	regen_p.add_argument("-n", "--name", required=True, help="The model ID/name.")
	regen_p.add_argument("-t", "--type", required=True, choices=["cloud", "local"], help="Model source type.")

	subparsers.add_parser("sync", help="Sync the local output directory with the HF bucket.")
	return parser

	def get_md5(file_path):
	"""Generate a quick hash of the file content."""
	hasher = hashlib.md5()
	with open(file_path, 'rb') as f:
	for chunk in iter(lambda: f.read(4096), b""):
	hasher.update(chunk)
	return hasher.hexdigest()
	def clean_local_renames_before_sync(local_dir):
	"""Finds duplicate files by content and removes the older-named version."""
	seen_hashes = {}
	for root, _, files in os.walk(local_dir):
	for file in files:
	path = os.path.join(root, file)
	f_hash = get_md5(path)

	if f_hash in seen_hashes:
	# We found a rename/duplicate! Delete the ghost/old copy
	print(f"⚠️ Detected rename conflict. Removing local ghost file: {path}")
	os.remove(path)
	else:
	seen_hashes[f_hash] = path

	def sync(output_dir):
	clean_local_renames_before_sync(output_dir)

	sync_bucket("hf://buckets/inferenceport-ai/model-info", output_dir, delete=False)

	sync_bucket(output_dir, "hf://buckets/inferenceport-ai/model-info", delete=False)

	def main():
	parser = build_arg_parser()
	args = parser.parse_args()

	api_key = args.api_key or DEFAULT_API_KEY
	output_dir = args.output_dir

	if api_key in ("YOUR_POLLINATIONS_API_KEY_HERE", "", None):
	console.print(
	"[bold yellow]Warning:[/bold yellow] API_KEY is still a placeholder. "
	"Pass --api-key or set DEFAULT_API_KEY before running.\n"
	)

	os.makedirs(output_dir, exist_ok=True)
	os.makedirs(os.path.join(output_dir, "compressed"), exist_ok=True)

	command = args.command or "scan"

	if command == "scan":
	cmd_scan(output_dir, api_key)
	elif command == "create":
	cmd_create(args.name, args.type, output_dir, api_key)
	elif command == "regenerate":
	cmd_regenerate(args.name, args.type, output_dir, api_key)
	elif command == "sync":
	sync(output_dir)
	else:
	parser.print_help()


	if __name__ == "__main__":
	main()