config.py

"""
Configuration file for Pixagram AI Pixel Art Generator
"""
import os
import torch

# Device configuration
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

# Model configuration
MODEL_REPO = "primerz/pixagram"
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN", None)  # Get token from environment

# Model files
MODEL_FILES = {
    "checkpoint": "horizon.safetensors",
    "lora": "retroart.safetensors",
    "vae": "pixelate.safetensors"
}

# LORA configuration
TRIGGER_WORD = "p1x3l4rt, pixel art"

# Face detection configuration
FACE_DETECTION_CONFIG = {
    "model_name": "antelopev2",
    "det_size": (640, 640),
    "ctx_id": 0
}

# Recommended resolutions
RECOMMENDED_SIZES = [
    (896, 1152),   # Portrait
    (1152, 896),   # Landscape
    (832, 1216),   # Tall portrait
    (1216, 832),   # Wide landscape
    (1024, 1024)   # Square
]

# Default generation parameters
DEFAULT_PARAMS = {
    "num_inference_steps": 12,
    "guidance_scale": 1.3,
    "strength": 0.50,
    "depth_control_scale": 0.75,
    "identity_control_scale": 0.85,
    "lora_scale": 1.0,
    "identity_preservation": 1.2,
    "enable_color_matching": False,
    "consistency_mode": True,
    "seed": -1
}

# Optimized preset configurations
PRESETS = {
    "Ultra Fidelity": {
        "strength": 0.40,
        "guidance_scale": 1.15,
        "identity_preservation": 1.8,
        "lora_scale": 0.8,
        "depth_control_scale": 0.65,
        "identity_control_scale": 0.95,
        "description": "Maximum face preservation - 96-98% similarity (Level 3)"
    },
    "Premium Portrait": {
        "strength": 0.52,
        "guidance_scale": 1.3,
        "identity_preservation": 1.35,
        "lora_scale": 1.1,
        "depth_control_scale": 0.75,
        "identity_control_scale": 0.85,
        "description": "Optimized balanced - strong pixel art + excellent face (92-94% similarity)"
    },
    "Balanced Portrait": {
        "strength": 0.50,
        "guidance_scale": 1.3,
        "identity_preservation": 1.2,
        "lora_scale": 1.0,
        "depth_control_scale": 0.75,
        "identity_control_scale": 0.85,
        "description": "Good balance between fidelity and style - 90-93% similarity"
    },
    "Artistic Excellence": {
        "strength": 0.58,
        "guidance_scale": 1.4,
        "identity_preservation": 1.2,
        "lora_scale": 1.2,
        "depth_control_scale": 0.78,
        "identity_control_scale": 0.75,
        "description": "Creative with strong likeness - 88-91% similarity"
    },
    "Style Focus": {
        "strength": 0.68,
        "guidance_scale": 1.5,
        "identity_preservation": 0.9,
        "lora_scale": 1.4,
        "depth_control_scale": 0.82,
        "identity_control_scale": 0.65,
        "description": "Maximum pixel art style - 83-87% similarity"
    },
    "Subtle Enhancement": {
        "strength": 0.38,
        "guidance_scale": 1.1,
        "identity_preservation": 1.9,
        "lora_scale": 0.75,
        "depth_control_scale": 0.60,
        "identity_control_scale": 0.98,
        "description": "Minimal transformation, photo-realistic - 97-99% similarity"
    }
}

# Multi-scale face processing
MULTI_SCALE_FACTORS = [0.75, 1.0, 1.25]

# Adaptive parameter adjustment thresholds
ADAPTIVE_THRESHOLDS = {
    "small_face_size": 50000,
    "low_confidence": 0.8,
    "profile_angle": 20
}

# Adaptive parameter sets
ADAPTIVE_PARAMS = {
    "small_face": {
        "identity_preservation": 1.8,
        "identity_control_scale": 0.95,
        "guidance_scale": 1.2,
        "lora_scale": 0.8,
        "reason": "Small face detected - boosting preservation"
    },
    "low_confidence": {
        "identity_preservation": 1.6,
        "identity_control_scale": 0.9,
        "guidance_scale": 1.3,
        "lora_scale": 0.85,
        "reason": "Low confidence - increasing identity weight"
    },
    "profile_view": {
        "identity_preservation": 1.7,
        "identity_control_scale": 0.95,
        "guidance_scale": 1.2,
        "lora_scale": 0.85,
        "reason": "Profile view - enhancing preservation"
    }
}

# Caption generation settings
CAPTION_CONFIG = {
    "max_length": 20,
    "num_beams": 4
}

# Color matching settings
COLOR_MATCH_CONFIG = {
    "lab_lightness_blend": 0.15,  # 15% adjustment to L channel
    "lab_color_blend_preserved": 0.05,  # 5% adjustment with saturation preservation
    "lab_color_blend_full": 0.20,  # 20% adjustment without preservation
    "saturation_boost": 1.05,  # Minimal saturation boost
    "gaussian_blur_kernel": (51, 51),
    "gaussian_blur_sigma": 20
}

# Face mask settings
FACE_MASK_CONFIG = {
    "padding": 0.1,  # 10% padding around face
    "feather": 30  # Blur radius for soft edges
}

# Model download retry settings
DOWNLOAD_CONFIG = {
    "max_retries": 3,
    "retry_delay": 2  # seconds
}

# Age brackets for demographic detection
AGE_BRACKETS = [
    (0, 18, "young"),
    (18, 30, "young adult"),
    (30, 50, "middle-aged"),
    (50, 150, "mature")
]

# CLIP skip setting
CLIP_SKIP = 2

# Identity boost multiplier
IDENTITY_BOOST_MULTIPLIER = 1.15

print(f"[CONFIG] Loaded configuration")
print(f"  Device: {device}")
print(f"  Dtype: {dtype}")
print(f"  Model Repo: {MODEL_REPO}")
print(f"  HuggingFace Token: {'Set' if HUGGINGFACE_TOKEN else 'Not set (using IP-based access)'}")