config.py

"""
Configuration file for Phramer AI
By Pariente AI, for MIA TV Series

Multimodal tool with BAGEL integration and professional photographic prompt optimization
"""

import os
import torch
from typing import Dict, Any

# Application Configuration
APP_CONFIG = {
    "title": "Phramer AI",
    "description": "Multimodal tool that reads images and turns them into refined, photo-realistic prompts. Ready for Midjourney, Flux or any generative engine.",
    "version": "2.0.0",
    "author": "Pariente AI",
    "project": "MIA TV Series",
    "tagline": "By Pariente AI, for MIA TV Series",
    "logline": "Phramer AI is a multimodal tool that reads an image and turns it into a refined, photo-realistic prompt. Ready for Midjourney, Flux or any generative engine."
}

# BAGEL Model Configuration
BAGEL_CONFIG = {
    "model_repo": "ByteDance-Seed/BAGEL-7B-MoT",
    "local_model_path": "./model",
    "cache_dir": "./model/cache",
    "download_patterns": ["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
    
    # Model parameters
    "dtype": torch.bfloat16,
    "device_map_strategy": "auto",
    "max_memory_per_gpu": "80GiB",
    "offload_buffers": True,
    "force_hooks": True,
    
    # Image processing
    "vae_transform_size": (1024, 512, 16),
    "vit_transform_size": (980, 224, 14),
    
    # Inference parameters
    "max_new_tokens": 512,
    "temperature": 0.7,
    "top_p": 0.9,
    "do_sample": True
}

# Device Configuration for ZeroGPU
def get_device_config() -> Dict[str, Any]:
    """Determine optimal device configuration for BAGEL"""
    device_config = {
        "device": "cpu",
        "use_gpu": False,
        "gpu_count": 0,
        "memory_efficient": True
    }
    
    if torch.cuda.is_available():
        gpu_count = torch.cuda.device_count()
        device_config.update({
            "device": "cuda",
            "use_gpu": True,
            "gpu_count": gpu_count,
            "gpu_memory_gb": torch.cuda.get_device_properties(0).total_memory / 1e9,
            "multi_gpu": gpu_count > 1
        })
    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        device_config.update({
            "device": "mps",
            "use_gpu": True,
            "gpu_count": 1
        })
    
    return device_config

# BAGEL Device Mapping Configuration
def get_bagel_device_map(gpu_count: int) -> Dict[str, str]:
    """Configure device mapping for BAGEL model"""
    # Same device modules that need to be on the same GPU
    same_device_modules = [
        'language_model.model.embed_tokens',
        'time_embedder',
        'latent_pos_embed',
        'vae2llm',
        'llm2vae',
        'connector',
        'vit_pos_embed'
    ]
    
    device_map = {}
    
    if gpu_count == 1:
        # Single GPU configuration
        for module in same_device_modules:
            device_map[module] = "cuda:0"
    else:
        # Multi-GPU configuration - keep critical modules on same device
        first_device = "cuda:0"
        for module in same_device_modules:
            device_map[module] = first_device
    
    return device_map

# Processing Configuration
PROCESSING_CONFIG = {
    "max_image_size": 1024,
    "image_quality": 95,
    "supported_formats": [".jpg", ".jpeg", ".png", ".webp"],
    "batch_size": 1,
    "timeout_seconds": 120  # Increased for BAGEL processing
}

# Prompt Optimization Rules for Multi-Engine Compatibility
FLUX_RULES = {
    "remove_patterns": [
        r',\s*trending on artstation',
        r',\s*trending on [^,]+',
        r',\s*\d+k\s*',
        r',\s*\d+k resolution',
        r',\s*artstation',
        r',\s*concept art',
        r',\s*digital art',
        r',\s*by greg rutkowski',
    ],
    
    "camera_configs": {
        "portrait": ", Shot on Hasselblad X2D 100C, 90mm f/2.5 lens at f/2.8, professional portrait photography",
        "landscape": ", Shot on Phase One XT, 40mm f/4 lens at f/8, epic landscape photography", 
        "street": ", Shot on Leica M11, 35mm f/1.4 lens at f/2.8, documentary street photography",
        "cinematic": ", Shot on ARRI Alexa LF, 35mm anamorphic lens, cinematic lighting",
        "architectural": ", Shot on Canon EOS R5, 24-70mm f/2.8 lens at f/8, architectural photography",
        "default": ", Shot on Phase One XF IQ4, 80mm f/2.8 lens at f/4, professional photography"
    },
    
    "lighting_enhancements": {
        "dramatic": ", dramatic cinematic lighting",
        "portrait": ", professional studio lighting with subtle rim light",
        "cinematic": ", moody cinematic lighting with practical lights",
        "natural": ", natural lighting with soft shadows",
        "default": ", masterful natural lighting"
    },
    
    "style_enhancements": {
        "photorealistic": ", photorealistic, ultra-detailed",
        "cinematic": ", cinematic composition, film grain",
        "commercial": ", commercial photography style",
        "editorial": ", editorial photography style",
        "fine_art": ", fine art photography"
    }
}

# Enhanced Scoring Configuration with Professional Photography Criteria
SCORING_CONFIG = {
    "max_score": 100,
    "score_weights": {
        "prompt_quality": 0.25,
        "technical_details": 0.25,
        "professional_photography": 0.25,
        "multi_engine_optimization": 0.25
    },
    
    "grade_thresholds": {
        95: {"grade": "LEGENDARY", "color": "#059669"},
        90: {"grade": "EXCELLENT", "color": "#10b981"},
        85: {"grade": "VERY GOOD", "color": "#22c55e"},
        75: {"grade": "GOOD", "color": "#84cc16"},
        65: {"grade": "FAIR", "color": "#f59e0b"},
        50: {"grade": "NEEDS WORK", "color": "#f97316"},
        0: {"grade": "POOR", "color": "#ef4444"}
    },
    
    "professional_criteria": {
        "camera_equipment": ["Canon", "Sony", "Leica", "Hasselblad", "Phase One", "ARRI"],
        "lens_specifications": ["f/", "mm", "anamorphic", "telephoto", "wide-angle"],
        "lighting_techniques": ["cinematic", "dramatic", "natural", "studio", "rim light"],
        "composition_rules": ["rule of thirds", "leading lines", "depth of field", "bokeh"],
        "technical_settings": ["aperture", "ISO", "shutter speed", "exposure"]
    }
}

# Environment Configuration
ENVIRONMENT = {
    "is_spaces": os.getenv("SPACE_ID") is not None,
    "is_local": os.getenv("SPACE_ID") is None,
    "log_level": os.getenv("LOG_LEVEL", "INFO"),
    "debug_mode": os.getenv("DEBUG", "false").lower() == "true",
    "space_id": os.getenv("SPACE_ID", ""),
    "space_author": os.getenv("SPACE_AUTHOR_NAME", "")
}

# Enhanced BAGEL Prompts for Professional Analysis
BAGEL_PROMPTS = {
    "multimodal_analysis": """Analyze this image for professional prompt generation. Provide exactly two sections:

1. DESCRIPTION: Create a detailed, flowing paragraph describing the image, including:
- Image type (photograph, illustration, artwork, scene)
- Subject matter and composition elements
- Color palette, lighting conditions, and mood
- Visual style, artistic elements, and photographic techniques
- Any cinematic or dramatic qualities

2. CAMERA_SETUP: Recommend professional camera and lens setup based on scene analysis:
- For portraits: High-end camera with portrait lens (85mm-135mm f/1.4-f/2.8)
- For landscapes: Medium format with wide lens (24-40mm f/4-f/8)
- For street/documentary: Compact system with standard lens (35mm-50mm f/1.4-f/2.8)
- For cinematic scenes: Cinema camera with appropriate lens for mood
- Include specific aperture, focal length, and shooting considerations

Focus on creating prompts suitable for multiple generative engines (Flux, Midjourney, etc.).""",

    "flux_optimization": """Analyze this image for FLUX prompt generation with professional photography principles:

1. DESCRIPTION: Detailed scene description focusing on:
- Visual composition and subject placement
- Lighting quality, direction, and mood
- Color relationships and tonal values
- Artistic style and photographic approach
- Technical qualities that enhance realism

2. CAMERA_SETUP: Professional equipment recommendation:
- Camera body suitable for the scene type
- Lens choice with specific focal length and aperture
- Lighting setup considerations
- Technical settings for optimal results

Generate content optimized for photorealistic output.""",

    "cinematic_analysis": """Analyze this image for cinematic prompt creation:

1. DESCRIPTION: Focus on cinematic elements:
- Scene composition and framing
- Lighting mood and dramatic elements
- Color grading and visual atmosphere
- Character positioning and environmental context
- Storytelling visual cues

2. CAMERA_SETUP: Cinema-grade equipment:
- Professional cinema camera recommendation
- Lens choice for cinematic effect
- Lighting setup for mood and story
- Technical considerations for film-quality output

Optimize for high-end generative engines with cinematic capabilities."""
}

# Professional Photography Integration Settings
PROFESSIONAL_PHOTOGRAPHY_CONFIG = {
    "enable_expert_analysis": True,
    "knowledge_base_integration": True,
    "technical_enhancement": True,
    "composition_guidance": True,
    
    "scene_detection_keywords": {
        "portrait": ["person", "face", "portrait", "headshot", "model"],
        "landscape": ["landscape", "nature", "mountain", "sky", "horizon"],
        "street": ["street", "urban", "city", "documentary", "candid"],
        "architectural": ["building", "architecture", "structure", "interior"],
        "cinematic": ["film", "movie", "cinematic", "dramatic", "story"],
        "commercial": ["product", "commercial", "advertising", "brand"]
    },
    
    "enhancement_priorities": [
        "technical_accuracy",
        "professional_terminology", 
        "equipment_specifications",
        "lighting_description",
        "composition_analysis"
    ]
}

# Flash Attention Installation Command
FLASH_ATTN_INSTALL = {
    "command": "pip install flash-attn --no-build-isolation",
    "env": {"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
    "shell": True
}

# Export main configurations
__all__ = [
    "APP_CONFIG",
    "BAGEL_CONFIG",
    "get_device_config",
    "get_bagel_device_map",
    "PROCESSING_CONFIG",
    "FLUX_RULES",
    "SCORING_CONFIG",
    "ENVIRONMENT",
    "BAGEL_PROMPTS",
    "PROFESSIONAL_PHOTOGRAPHY_CONFIG",
    "FLASH_ATTN_INSTALL"
]