config.py

"""
Configuration file for Frame 0 Laboratory for MIA
BAGEL 7B integration with FLUX prompt optimization
"""

import os
import torch
from typing import Dict, Any

# Application Configuration
APP_CONFIG = {
    "title": "Frame 0 Laboratory for MIA",
    "description": "Advanced image analysis with BAGEL 7B and FLUX prompt optimization",
    "version": "2.0.0",
    "author": "Frame 0 Laboratory for MIA"
}

# 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
}

# FLUX Prompt Rules
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",
        "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",
        "default": ", masterful natural lighting"
    }
}

# Scoring Configuration
SCORING_CONFIG = {
    "max_score": 100,
    "score_weights": {
        "prompt_quality": 0.3,
        "technical_details": 0.25,
        "artistic_value": 0.25,
        "flux_optimization": 0.2
    },
    
    "grade_thresholds": {
        95: {"grade": "LEGENDARY", "color": "#059669"},
        90: {"grade": "EXCELLENT", "color": "#10b981"},
        80: {"grade": "VERY GOOD", "color": "#22c55e"},
        70: {"grade": "GOOD", "color": "#f59e0b"},
        60: {"grade": "FAIR", "color": "#f97316"},
        0: {"grade": "NEEDS WORK", "color": "#ef4444"}
    }
}

# 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", "")
}

# BAGEL Inference Prompts
BAGEL_PROMPTS = {
    "image_analysis": "Describe this image in detail, including objects, people, setting, mood, and visual elements:",
    "flux_prompt": "Generate a detailed FLUX prompt for this image, focusing on photographic and artistic elements:",
    "detailed_description": "Provide a comprehensive analysis of this image including composition, lighting, colors, and artistic style:",
}

# 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",
    "FLASH_ATTN_INSTALL"
]