Update config.py

config.py

@@ -1,6 +1,6 @@
 """
-Configuration file for FLUX Prompt Optimizer
-Clean, simple, and focused configuration
+Configuration file for Frame 0 Laboratory for MIA
+BAGEL 7B integration with FLUX prompt optimization
 """
 
 import os
@@ -9,67 +9,103 @@ from typing import Dict, Any
 
 # Application Configuration
 APP_CONFIG = {
-    "title": "🚀 FLUX Prompt Optimizer",
-    "description": "Advanced image analysis and Flux prompt optimization",
+    "title": "Frame 0 Laboratory for MIA",
+    "description": "Advanced image analysis with BAGEL 7B and FLUX prompt optimization",
     "version": "2.0.0",
-    "author": "Pariente AI Research"
+    "author": "Frame 0 Laboratory for MIA"
 }
 
-# Model Configuration
-MODEL_CONFIG = {
-    # Primary analysis model - choose one
-    "primary_model": "florence2",  # or "bagel" 
+# 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"],
     
-    # Florence-2 settings
-    "florence2": {
-        "model_id": "microsoft/Florence-2-base",
-        "torch_dtype": torch.float16,
-        "device_map": "auto",
-        "trust_remote_code": True,
-        "max_new_tokens": 1024
-    },
+    # Model parameters
+    "dtype": torch.bfloat16,
+    "device_map_strategy": "auto",
+    "max_memory_per_gpu": "80GiB",
+    "offload_buffers": True,
+    "force_hooks": True,
     
-    # Bagel-7B settings (via API)
-    "bagel": {
-        "api_url": "https://huggingface.co/spaces/Malaji71/Bagel-7B-Demo",
-        "timeout": 30,
-        "max_retries": 3
-    }
+    # 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
+# Device Configuration for ZeroGPU
 def get_device_config() -> Dict[str, Any]:
-    """Determine optimal device configuration"""
+    """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_memory_gb": torch.cuda.get_device_properties(0).total_memory / 1e9
+            "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
+            "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": 60
+    "timeout_seconds": 120  # Increased for BAGEL processing
 }
 
-# Flux Prompt Rules
+# FLUX Prompt Rules
 FLUX_RULES = {
     "remove_patterns": [
         r',\s*trending on artstation',
@@ -121,16 +157,35 @@ 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"
+    "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",
-    "MODEL_CONFIG", 
+    "BAGEL_CONFIG",
     "get_device_config",
+    "get_bagel_device_map",
     "PROCESSING_CONFIG",
     "FLUX_RULES",
     "SCORING_CONFIG",
-    "ENVIRONMENT"
+    "ENVIRONMENT",
+    "BAGEL_PROMPTS",
+    "FLASH_ATTN_INSTALL"
 ]