slahiriandClaude Opus 4.5
c90fac7039
- Consolidate 3 text boxes (Enhance/Replace/Add) into single Instructions field - Add VRAM clearing and garbage collection at start of all prompt nodes - Add smart image compression with file size limits (500KB local, 1MB API) - Add compressed_image output for preview in Generator - Add toggle to disable negative prompt generation (default OFF) - Add unload_model() to LocalModelClient to free VRAM after use - Replace detail_level with analysis_mode String input for wire compatibility - Rename outputs to out_prompt_length/out_analysis_mode to avoid ID conflicts - Strengthen clothing color accuracy in prompts - Remove deprecated sid_zimage_optimizer.py 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
763 lines
23 KiB
Python
763 lines
23 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
ComfyUI-AI-Photography-Toolkit - Z-Image Utilities
|
|
|
|
Utility functions for config loading, caching, and output cleaning.
|
|
|
|
Author: Siddhartha Lahiri
|
|
Email: siddhartha.lahiri@gmail.com
|
|
License: MIT
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
import json
|
|
import hashlib
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
import yaml
|
|
import numpy as np
|
|
|
|
|
|
# =============================================================================
|
|
# CONFIG LOADING
|
|
# =============================================================================
|
|
|
|
_CONFIG_CACHE: dict[str, Any] = {}
|
|
|
|
|
|
def get_config_path() -> Path:
|
|
"""Get the path to the config directory."""
|
|
return Path(__file__).parent.parent / "config"
|
|
|
|
|
|
def load_zimage_config(force_reload: bool = False) -> dict[str, Any]:
|
|
"""
|
|
Load the Z-Image generator configuration from YAML.
|
|
|
|
Args:
|
|
force_reload: Force reload from disk even if cached
|
|
|
|
Returns:
|
|
Configuration dictionary
|
|
"""
|
|
global _CONFIG_CACHE
|
|
|
|
cache_key = "zimage_generator"
|
|
|
|
if not force_reload and cache_key in _CONFIG_CACHE:
|
|
return _CONFIG_CACHE[cache_key]
|
|
|
|
config_path = get_config_path() / "zimage_generator_config.yaml"
|
|
|
|
if not config_path.exists():
|
|
raise FileNotFoundError(f"Config file not found: {config_path}")
|
|
|
|
with open(config_path, "r", encoding="utf-8") as f:
|
|
config = yaml.safe_load(f)
|
|
|
|
_CONFIG_CACHE[cache_key] = config
|
|
return config
|
|
|
|
|
|
def get_shot_framings() -> dict[str, dict]:
|
|
"""Get all shot framing definitions."""
|
|
config = load_zimage_config()
|
|
return config.get("shot_framings", {})
|
|
|
|
|
|
def get_photography_genres() -> dict[str, dict]:
|
|
"""Get all photography genre definitions, flattened."""
|
|
config = load_zimage_config()
|
|
genres = config.get("photography_genres", {})
|
|
|
|
# Flatten the nested structure
|
|
flattened = {}
|
|
for category, category_genres in genres.items():
|
|
for code, genre_data in category_genres.items():
|
|
genre_data["category"] = category
|
|
flattened[code] = genre_data
|
|
|
|
return flattened
|
|
|
|
|
|
def get_genre_options() -> list[str]:
|
|
"""Get list of genre codes for UI dropdown."""
|
|
genres = get_photography_genres()
|
|
return list(genres.keys())
|
|
|
|
|
|
def get_attribute_schema(attribute_name: str) -> dict[str, Any]:
|
|
"""Get the schema for a specific attribute."""
|
|
config = load_zimage_config()
|
|
attributes = config.get("attributes", {})
|
|
return attributes.get(attribute_name, {})
|
|
|
|
|
|
def get_content_detail_schema(level: str) -> dict[str, Any]:
|
|
"""Get the attribute schema for a specific content detail level."""
|
|
config = load_zimage_config()
|
|
levels = config.get("content_detail_levels", {})
|
|
return levels.get(level, levels.get("standard", {}))
|
|
|
|
|
|
def get_prompt_template(stage: str) -> str:
|
|
"""Get the prompt template for a specific stage."""
|
|
config = load_zimage_config()
|
|
templates = config.get("prompt_templates", {})
|
|
stage_template = templates.get(stage, {})
|
|
return stage_template.get("system", "")
|
|
|
|
|
|
def get_focus_override_config(override_name: str) -> Optional[dict]:
|
|
"""Get configuration for a focus override."""
|
|
config = load_zimage_config()
|
|
overrides = config.get("focus_overrides", {})
|
|
return overrides.get(override_name)
|
|
|
|
|
|
def get_detail_level_config(level: str) -> dict[str, Any]:
|
|
"""Get configuration for a detail level."""
|
|
config = load_zimage_config()
|
|
levels = config.get("detail_levels", {})
|
|
return levels.get(level, levels.get("Standard", {}))
|
|
|
|
|
|
def get_zimage_settings() -> dict[str, Any]:
|
|
"""Get Z-Image recommended settings."""
|
|
config = load_zimage_config()
|
|
return config.get("zimage_settings", {})
|
|
|
|
|
|
# =============================================================================
|
|
# OUTPUT CLEANING
|
|
# =============================================================================
|
|
|
|
def clean_zimage_output(text: str, max_length: int = 2000) -> str:
|
|
"""
|
|
Clean LLM output for Z-Image compatibility.
|
|
|
|
Args:
|
|
text: Raw LLM output
|
|
max_length: Maximum character length (0 = unlimited)
|
|
|
|
Returns:
|
|
Cleaned text suitable for Z-Image
|
|
"""
|
|
if not text:
|
|
return ""
|
|
|
|
# Remove thinking tags
|
|
text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
|
|
text = re.sub(r"<thinking>.*?</thinking>", "", text, flags=re.DOTALL)
|
|
|
|
# Remove markdown code blocks
|
|
text = re.sub(r"```[\s\S]*?```", "", text)
|
|
text = re.sub(r"`[^`]+`", "", text)
|
|
|
|
# Remove common prefixes
|
|
prefixes_to_remove = [
|
|
r"^Here['']?s? (?:is )?(?:the |a |an )?(?:enhanced |detailed |generated )?prompt[:\s]*",
|
|
r"^(?:The )?(?:enhanced |detailed |generated )?prompt[:\s]*",
|
|
r"^Sure[,!]? (?:here['']?s? )?(?:the |a )?prompt[:\s]*",
|
|
r"^I['']?ve (?:generated|created|written)[:\s]*",
|
|
]
|
|
for prefix in prefixes_to_remove:
|
|
text = re.sub(prefix, "", text, flags=re.IGNORECASE)
|
|
|
|
# Remove meta-tags that Z-Image doesn't need
|
|
meta_tags = [
|
|
r"\b8[kK]\b",
|
|
r"\bmasterpiece\b",
|
|
r"\bbest quality\b",
|
|
r"\bhigh quality\b",
|
|
r"\bultra quality\b",
|
|
r"\b4[kK]\b",
|
|
r"\bUHD\b",
|
|
r"\bhighly detailed\b",
|
|
r"\bsuper detailed\b",
|
|
]
|
|
for tag in meta_tags:
|
|
text = re.sub(tag + r",?\s*", "", text, flags=re.IGNORECASE)
|
|
|
|
# Clean up whitespace
|
|
text = re.sub(r"\s+", " ", text)
|
|
text = text.strip()
|
|
|
|
# Remove leading/trailing commas
|
|
text = text.strip(",").strip()
|
|
|
|
# Clean up double commas
|
|
text = re.sub(r",\s*,", ",", text)
|
|
|
|
# Truncate if needed (at sentence boundary)
|
|
if max_length > 0 and len(text) > max_length:
|
|
text = truncate_at_sentence(text, max_length)
|
|
|
|
return text
|
|
|
|
|
|
def truncate_at_sentence(text: str, max_length: int) -> str:
|
|
"""
|
|
Truncate text at the nearest sentence boundary before max_length.
|
|
|
|
Args:
|
|
text: Text to truncate
|
|
max_length: Maximum length
|
|
|
|
Returns:
|
|
Truncated text ending at a sentence boundary
|
|
"""
|
|
if len(text) <= max_length:
|
|
return text
|
|
|
|
# Find sentence boundaries before max_length
|
|
truncated = text[:max_length]
|
|
|
|
# Look for sentence-ending punctuation
|
|
last_period = truncated.rfind(". ")
|
|
last_exclaim = truncated.rfind("! ")
|
|
last_question = truncated.rfind("? ")
|
|
|
|
# Find the latest sentence boundary
|
|
boundary = max(last_period, last_exclaim, last_question)
|
|
|
|
if boundary > max_length * 0.5: # Only use if at least halfway through
|
|
return truncated[:boundary + 1].strip()
|
|
|
|
# Fallback: truncate at last space
|
|
last_space = truncated.rfind(" ")
|
|
if last_space > max_length * 0.7:
|
|
return truncated[:last_space].strip() + "..."
|
|
|
|
return truncated.strip() + "..."
|
|
|
|
|
|
# =============================================================================
|
|
# CACHING UTILITIES (Persistent Disk Cache)
|
|
# =============================================================================
|
|
|
|
# In-memory cache for fast repeated access within session
|
|
_MEMORY_CACHE: dict[str, Any] = {}
|
|
_MAX_MEMORY_CACHE_SIZE = 50
|
|
|
|
# Disk cache settings
|
|
_MAX_DISK_CACHE_SIZE = 500 # Max number of cached entries on disk
|
|
_CACHE_VERSION = "1" # Increment to invalidate old caches
|
|
|
|
|
|
def get_cache_dir() -> Path:
|
|
"""Get the cache directory path, creating it if needed."""
|
|
cache_dir = Path(__file__).parent.parent / "cache" / "prompts"
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
return cache_dir
|
|
|
|
|
|
def hash_image_tensor(image_tensor) -> str:
|
|
"""
|
|
Generate a hash from an image tensor for caching.
|
|
|
|
Args:
|
|
image_tensor: ComfyUI image tensor (B, H, W, C)
|
|
|
|
Returns:
|
|
Hex string hash of the image
|
|
"""
|
|
# Take first image if batch
|
|
if len(image_tensor.shape) == 4:
|
|
img_np = image_tensor[0].cpu().numpy()
|
|
else:
|
|
img_np = image_tensor.cpu().numpy()
|
|
|
|
# Downsample for faster hashing (every 4th pixel)
|
|
downsampled = img_np[::4, ::4, :]
|
|
|
|
# Create hash from bytes
|
|
img_bytes = downsampled.tobytes()
|
|
return hashlib.md5(img_bytes).hexdigest()
|
|
|
|
|
|
def get_cache_key(
|
|
image_hash: str,
|
|
seed: int,
|
|
**params
|
|
) -> str:
|
|
"""
|
|
Generate a cache key from image hash and parameters.
|
|
|
|
Args:
|
|
image_hash: Hash of the input image
|
|
seed: Random seed
|
|
**params: All other parameters that affect output
|
|
|
|
Returns:
|
|
Cache key string
|
|
"""
|
|
# Sort params for consistent ordering
|
|
sorted_params = sorted(params.items())
|
|
param_str = json.dumps(sorted_params, sort_keys=True)
|
|
param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8]
|
|
|
|
return f"v{_CACHE_VERSION}_{image_hash}_{seed}_{param_hash}"
|
|
|
|
|
|
def _get_cache_file_path(cache_key: str) -> Path:
|
|
"""Get the file path for a cache entry."""
|
|
# Use first 2 chars as subdirectory for better file distribution
|
|
subdir = cache_key[:2] if len(cache_key) >= 2 else "00"
|
|
cache_dir = get_cache_dir() / subdir
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
return cache_dir / f"{cache_key}.json"
|
|
|
|
|
|
def get_cached_output(cache_key: str) -> Optional[dict]:
|
|
"""
|
|
Get cached output if it exists (checks memory first, then disk).
|
|
|
|
Args:
|
|
cache_key: The cache key to look up
|
|
|
|
Returns:
|
|
Cached output dict or None if not found
|
|
"""
|
|
global _MEMORY_CACHE
|
|
|
|
# Check memory cache first (fast path)
|
|
if cache_key in _MEMORY_CACHE:
|
|
return _MEMORY_CACHE[cache_key]
|
|
|
|
# Check disk cache
|
|
cache_file = _get_cache_file_path(cache_key)
|
|
if cache_file.exists():
|
|
try:
|
|
with open(cache_file, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
# Validate cache version
|
|
if data.get("_cache_version") == _CACHE_VERSION:
|
|
output = data.get("output")
|
|
# Store in memory cache for faster subsequent access
|
|
_MEMORY_CACHE[cache_key] = output
|
|
return output
|
|
else:
|
|
# Old cache version, delete it
|
|
cache_file.unlink(missing_ok=True)
|
|
except (json.JSONDecodeError, IOError):
|
|
# Corrupted cache file, delete it
|
|
cache_file.unlink(missing_ok=True)
|
|
|
|
return None
|
|
|
|
|
|
def set_cached_output(cache_key: str, output: dict) -> None:
|
|
"""
|
|
Cache an output result (both memory and disk).
|
|
|
|
Args:
|
|
cache_key: The cache key
|
|
output: The output dict to cache
|
|
"""
|
|
global _MEMORY_CACHE
|
|
|
|
# Store in memory cache
|
|
if len(_MEMORY_CACHE) >= _MAX_MEMORY_CACHE_SIZE:
|
|
# Remove oldest entries (first 20%)
|
|
keys_to_remove = list(_MEMORY_CACHE.keys())[:_MAX_MEMORY_CACHE_SIZE // 5]
|
|
for key in keys_to_remove:
|
|
del _MEMORY_CACHE[key]
|
|
|
|
_MEMORY_CACHE[cache_key] = output
|
|
|
|
# Store on disk
|
|
cache_file = _get_cache_file_path(cache_key)
|
|
try:
|
|
cache_data = {
|
|
"_cache_version": _CACHE_VERSION,
|
|
"_created": str(Path(cache_file).stat().st_mtime if cache_file.exists() else "new"),
|
|
"output": output
|
|
}
|
|
with open(cache_file, "w", encoding="utf-8") as f:
|
|
json.dump(cache_data, f, indent=2)
|
|
|
|
# Cleanup old cache files if needed
|
|
_cleanup_disk_cache()
|
|
except IOError as e:
|
|
# Log but don't fail if disk cache write fails
|
|
print(f"Warning: Failed to write cache file: {e}")
|
|
|
|
|
|
def _cleanup_disk_cache() -> None:
|
|
"""Remove oldest cache files if over limit."""
|
|
cache_dir = get_cache_dir()
|
|
|
|
# Get all cache files with their modification times
|
|
cache_files = []
|
|
for subdir in cache_dir.iterdir():
|
|
if subdir.is_dir():
|
|
for cache_file in subdir.glob("*.json"):
|
|
try:
|
|
mtime = cache_file.stat().st_mtime
|
|
cache_files.append((cache_file, mtime))
|
|
except OSError:
|
|
pass
|
|
|
|
# If over limit, remove oldest files
|
|
if len(cache_files) > _MAX_DISK_CACHE_SIZE:
|
|
# Sort by modification time (oldest first)
|
|
cache_files.sort(key=lambda x: x[1])
|
|
|
|
# Remove oldest 20%
|
|
files_to_remove = cache_files[:len(cache_files) // 5]
|
|
for cache_file, _ in files_to_remove:
|
|
try:
|
|
cache_file.unlink(missing_ok=True)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def clear_cache(memory_only: bool = False) -> int:
|
|
"""
|
|
Clear the output cache.
|
|
|
|
Args:
|
|
memory_only: If True, only clear memory cache. If False, clear both.
|
|
|
|
Returns:
|
|
Number of entries cleared
|
|
"""
|
|
global _MEMORY_CACHE
|
|
|
|
count = len(_MEMORY_CACHE)
|
|
_MEMORY_CACHE = {}
|
|
|
|
if not memory_only:
|
|
cache_dir = get_cache_dir()
|
|
if cache_dir.exists():
|
|
for subdir in cache_dir.iterdir():
|
|
if subdir.is_dir():
|
|
for cache_file in subdir.glob("*.json"):
|
|
try:
|
|
cache_file.unlink()
|
|
count += 1
|
|
except OSError:
|
|
pass
|
|
|
|
return count
|
|
|
|
|
|
def get_cache_stats() -> dict[str, Any]:
|
|
"""Get cache statistics."""
|
|
cache_dir = get_cache_dir()
|
|
|
|
disk_count = 0
|
|
disk_size = 0
|
|
|
|
if cache_dir.exists():
|
|
for subdir in cache_dir.iterdir():
|
|
if subdir.is_dir():
|
|
for cache_file in subdir.glob("*.json"):
|
|
try:
|
|
disk_count += 1
|
|
disk_size += cache_file.stat().st_size
|
|
except OSError:
|
|
pass
|
|
|
|
return {
|
|
"memory_entries": len(_MEMORY_CACHE),
|
|
"memory_max": _MAX_MEMORY_CACHE_SIZE,
|
|
"disk_entries": disk_count,
|
|
"disk_max": _MAX_DISK_CACHE_SIZE,
|
|
"disk_size_mb": round(disk_size / (1024 * 1024), 2),
|
|
"cache_dir": str(cache_dir),
|
|
}
|
|
|
|
|
|
# =============================================================================
|
|
# IMAGE UTILITIES
|
|
# =============================================================================
|
|
|
|
def get_image_metadata(image_tensor) -> dict[str, Any]:
|
|
"""
|
|
Extract metadata from an image tensor.
|
|
|
|
Args:
|
|
image_tensor: ComfyUI image tensor (B, H, W, C)
|
|
|
|
Returns:
|
|
Dictionary with image metadata
|
|
"""
|
|
# Get dimensions
|
|
if len(image_tensor.shape) == 4:
|
|
batch, height, width, channels = image_tensor.shape
|
|
else:
|
|
height, width, channels = image_tensor.shape
|
|
batch = 1
|
|
|
|
# Calculate aspect ratio
|
|
aspect_decimal = width / height
|
|
|
|
# Determine orientation
|
|
if aspect_decimal > 1.05:
|
|
orientation = "landscape"
|
|
elif aspect_decimal < 0.95:
|
|
orientation = "portrait"
|
|
else:
|
|
orientation = "square"
|
|
|
|
# Find common aspect ratio
|
|
aspect_ratio = find_aspect_ratio(width, height)
|
|
|
|
# Calculate megapixels
|
|
megapixels = (width * height) / 1_000_000
|
|
|
|
return {
|
|
"width": width,
|
|
"height": height,
|
|
"channels": channels,
|
|
"batch_size": batch,
|
|
"aspect_ratio": aspect_ratio,
|
|
"aspect_decimal": round(aspect_decimal, 2),
|
|
"orientation": orientation,
|
|
"megapixels": round(megapixels, 2),
|
|
}
|
|
|
|
|
|
def find_aspect_ratio(width: int, height: int) -> str:
|
|
"""Find the closest common aspect ratio."""
|
|
from math import gcd
|
|
|
|
divisor = gcd(width, height)
|
|
w = width // divisor
|
|
h = height // divisor
|
|
|
|
# Common aspect ratios to match
|
|
common_ratios = {
|
|
(1, 1): "1:1",
|
|
(4, 3): "4:3",
|
|
(3, 4): "3:4",
|
|
(16, 9): "16:9",
|
|
(9, 16): "9:16",
|
|
(3, 2): "3:2",
|
|
(2, 3): "2:3",
|
|
(4, 5): "4:5",
|
|
(5, 4): "5:4",
|
|
(21, 9): "21:9",
|
|
(9, 21): "9:21",
|
|
}
|
|
|
|
# Check for exact match
|
|
if (w, h) in common_ratios:
|
|
return common_ratios[(w, h)]
|
|
|
|
# Find closest match
|
|
target_ratio = width / height
|
|
closest = None
|
|
min_diff = float("inf")
|
|
|
|
for (rw, rh), name in common_ratios.items():
|
|
ratio = rw / rh
|
|
diff = abs(ratio - target_ratio)
|
|
if diff < min_diff:
|
|
min_diff = diff
|
|
closest = name
|
|
|
|
# If very close to a common ratio, use it
|
|
if min_diff < 0.05:
|
|
return closest
|
|
|
|
# Otherwise return simplified ratio
|
|
return f"{w}:{h}"
|
|
|
|
|
|
def get_zimage_recommendations(image_metadata: dict) -> dict[str, Any]:
|
|
"""
|
|
Generate Z-Image recommendations based on image metadata.
|
|
|
|
Args:
|
|
image_metadata: Output from get_image_metadata()
|
|
|
|
Returns:
|
|
Dictionary with Z-Image recommendations
|
|
"""
|
|
settings = get_zimage_settings()
|
|
width = image_metadata["width"]
|
|
height = image_metadata["height"]
|
|
orientation = image_metadata["orientation"]
|
|
|
|
# Get optimal resolution based on orientation
|
|
if orientation == "portrait":
|
|
optimal = settings["optimal_resolution"]["portrait"]
|
|
elif orientation == "landscape":
|
|
optimal = settings["optimal_resolution"]["landscape"]
|
|
else:
|
|
optimal = settings["optimal_resolution"]["native"]
|
|
|
|
optimal_w, optimal_h = optimal
|
|
|
|
# Check if resize is needed
|
|
resize_needed = (width != optimal_w or height != optimal_h)
|
|
|
|
# Determine resize direction
|
|
if resize_needed:
|
|
if width * height > optimal_w * optimal_h:
|
|
resize_direction = "downscale"
|
|
else:
|
|
resize_direction = "upscale"
|
|
else:
|
|
resize_direction = "none"
|
|
|
|
# Check aspect ratio compatibility
|
|
current_aspect = width / height
|
|
optimal_aspect = optimal_w / optimal_h
|
|
aspect_ratio_ok = abs(current_aspect - optimal_aspect) < 0.1
|
|
|
|
# Estimate quality based on resolution
|
|
total_pixels = width * height
|
|
if total_pixels >= 1_000_000:
|
|
quality = "high"
|
|
elif total_pixels >= 500_000:
|
|
quality = "medium"
|
|
else:
|
|
quality = "low"
|
|
|
|
return {
|
|
"optimal_resolution": optimal,
|
|
"current_compatible": not resize_needed,
|
|
"resize_needed": resize_needed,
|
|
"resize_direction": resize_direction,
|
|
"resize_method": "lanczos",
|
|
"aspect_ratio_ok": aspect_ratio_ok,
|
|
"nearest_native": f"{optimal_w}x{optimal_h} ({image_metadata['aspect_ratio']})",
|
|
"quality_estimate": quality,
|
|
"suggested_settings": {
|
|
"steps": settings["inference"]["steps"],
|
|
"guidance_scale": settings["inference"]["guidance_scale"],
|
|
"max_sequence_length": settings["inference"]["max_sequence_length"],
|
|
},
|
|
}
|
|
|
|
|
|
# =============================================================================
|
|
# ATTRIBUTE SCHEMA BUILDERS
|
|
# =============================================================================
|
|
|
|
def build_attribute_schema_for_scene(
|
|
shot_framing: str,
|
|
genre: str,
|
|
content_detail: str = "standard"
|
|
) -> dict[str, list[str]]:
|
|
"""
|
|
Build the attribute schema to extract based on scene classification.
|
|
|
|
Args:
|
|
shot_framing: Shot framing code (e.g., "CU", "FS")
|
|
genre: Photography genre code (e.g., "PRT", "FSH")
|
|
content_detail: Content detail level
|
|
|
|
Returns:
|
|
Dictionary mapping attribute categories to fields to extract
|
|
"""
|
|
config = load_zimage_config()
|
|
shot_framings = config.get("shot_framings", {})
|
|
|
|
# Get base attributes from shot framing
|
|
framing_config = shot_framings.get(shot_framing, {})
|
|
base_attributes = framing_config.get("attributes", [])
|
|
|
|
# Get content detail level config
|
|
detail_config = get_content_detail_schema(content_detail)
|
|
|
|
# Build schema based on scene type
|
|
schema = {}
|
|
|
|
# Always include subject for people photography
|
|
genres = get_photography_genres()
|
|
genre_config = genres.get(genre, {})
|
|
genre_category = genre_config.get("category", "")
|
|
|
|
if genre_category == "people":
|
|
schema["subject"] = ["gender", "age_range"]
|
|
|
|
# Add attributes based on shot framing
|
|
for attr in base_attributes:
|
|
if attr == "face":
|
|
schema["face"] = ["shape", "angle", "tilt"]
|
|
elif attr == "eyes":
|
|
schema["eyes"] = ["color", "shape", "state", "gaze", "expression", "makeup"]
|
|
elif attr == "skin":
|
|
schema["skin"] = ["tone", "undertone", "texture", "condition"]
|
|
elif attr == "lips":
|
|
schema["lips"] = ["shape", "color", "state"]
|
|
elif attr == "expression":
|
|
schema["expression"] = ["overall", "intensity"]
|
|
elif attr == "hair":
|
|
schema["hair"] = ["color", "length", "texture", "style", "state"]
|
|
elif attr == "body_upper":
|
|
schema["body_upper"] = detail_config.get("body", ["build", "posture"])
|
|
schema["clothing_upper"] = detail_config.get("clothing", ["type", "color"])
|
|
elif attr == "body" or attr == "clothing_full":
|
|
schema["body"] = detail_config.get("body", ["build", "posture"])
|
|
if content_detail in ["detailed", "explicit"]:
|
|
schema["body_exposure"] = ["chest", "stomach", "back", "thigh", "buttocks"]
|
|
schema["clothing_full"] = detail_config.get("clothing", ["type", "color"])
|
|
elif attr == "pose_upper":
|
|
schema["pose_upper"] = ["shoulders", "arms", "hands", "body_angle"]
|
|
elif attr == "pose_full":
|
|
schema["pose_full"] = ["stance", "weight", "legs", "hips", "dynamic"]
|
|
if content_detail in ["detailed", "explicit"]:
|
|
schema["pose_full"].append("sensuality")
|
|
elif attr == "environment":
|
|
schema["environment"] = ["setting_type", "location"]
|
|
elif attr == "lighting":
|
|
schema["lighting"] = ["type", "source", "direction", "quality", "color_temperature"]
|
|
elif attr == "accessories":
|
|
schema["accessories"] = ["necklace", "earrings", "glasses"]
|
|
|
|
# Always include lighting and background for most shots
|
|
if "lighting" not in schema and shot_framing not in ["ECU", "BCU"]:
|
|
schema["lighting"] = ["type", "direction", "quality"]
|
|
|
|
if "background" not in schema and shot_framing not in ["ECU"]:
|
|
schema["background"] = ["clarity", "complexity", "color"]
|
|
|
|
return schema
|
|
|
|
|
|
def format_attribute_schema_for_prompt(schema: dict[str, list[str]]) -> str:
|
|
"""
|
|
Format an attribute schema as text for inclusion in a prompt.
|
|
|
|
Args:
|
|
schema: Attribute schema from build_attribute_schema_for_scene()
|
|
|
|
Returns:
|
|
Formatted text describing the schema
|
|
"""
|
|
config = load_zimage_config()
|
|
all_attributes = config.get("attributes", {})
|
|
|
|
lines = ["Extract the following attributes:\n"]
|
|
|
|
for category, fields in schema.items():
|
|
lines.append(f"\n{category}:")
|
|
|
|
category_config = all_attributes.get(category, {})
|
|
|
|
for field in fields:
|
|
field_config = category_config.get(field, {})
|
|
|
|
if isinstance(field_config, dict):
|
|
options = field_config.get("options", [])
|
|
if options:
|
|
options_str = ", ".join(options[:5])
|
|
if len(options) > 5:
|
|
options_str += ", ..."
|
|
lines.append(f" - {field}: [{options_str}]")
|
|
else:
|
|
lines.append(f" - {field}: (describe)")
|
|
else:
|
|
lines.append(f" - {field}: (describe)")
|
|
|
|
return "\n".join(lines)
|