Update processing/ai_background.py

processing/ai_background.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """
-AI Background Generator Module - Updated to handle dependency conflicts
-Handles Stable Diffusion background generation with graceful fallbacks.
+AI Background Generator Module - Gradient-Only Version
+Handles background generation with smart gradient fallbacks when AI libraries conflict.
 """
 
 import os
@@ -18,70 +18,162 @@
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Try to import AI libraries with fallback
-DIFFUSERS_AVAILABLE = False
+# Safe imports only
 PIL_AVAILABLE = False
+TORCH_AVAILABLE = False
+DIFFUSERS_AVAILABLE = False
 
 try:
     from PIL import Image, ImageDraw, ImageFilter
     PIL_AVAILABLE = True
-    logger.info("PIL imported successfully")
+    logger.info("✅ PIL imported successfully")
 except ImportError as e:
-    logger.warning(f"PIL import failed: {e}")
+    logger.error(f"❌ PIL import failed: {e}")
 
+# Check PyTorch availability but don't import if problematic
 try:
     import torch
-    from diffusers import StableDiffusionPipeline
-    DIFFUSERS_AVAILABLE = True
-    logger.info("Diffusers imported successfully")
+    TORCH_AVAILABLE = True
+    logger.info("✅ PyTorch available")
+    
+    # Test for the specific custom_op issue
+    if hasattr(torch.library, 'custom_op'):
+        logger.info("✅ PyTorch custom_op available")
+    else:
+        logger.warning("⚠️  PyTorch custom_op not available - diffusers will likely fail")
+        
+except ImportError:
+    logger.warning("⚠️  PyTorch not available")
+except AttributeError:
+    logger.warning("⚠️  PyTorch version incompatible (missing torch.library)")
 except Exception as e:
-    logger.warning(f"Diffusers import failed: {e}")
+    logger.warning(f"⚠️  PyTorch check failed: {e}")
+
+# NEVER attempt to import diffusers if we detect the custom_op issue
+FORCE_GRADIENT_MODE = False
+
+if TORCH_AVAILABLE:
+    try:
+        # First, check if torch.library.custom_op exists
+        import torch
+        if not hasattr(torch.library, 'custom_op'):
+            logger.warning("🔄 Force enabling gradient-only mode due to torch.library.custom_op missing")
+            FORCE_GRADIENT_MODE = True
+        else:
+            # Only try diffusers if custom_op exists
+            try:
+                # Quick test import to see if diffusers will work
+                import importlib.util
+                spec = importlib.util.find_spec("diffusers")
+                if spec is not None:
+                    # Try a minimal import test
+                    from diffusers import __version__
+                    logger.info(f"✅ Diffusers {__version__} detected and compatible")
+                    DIFFUSERS_AVAILABLE = True
+                else:
+                    logger.info("ℹ️  Diffusers not installed")
+            except Exception as e:
+                if "custom_op" in str(e):
+                    logger.warning("🔄 Detected custom_op compatibility issue - using gradient-only mode")
+                    FORCE_GRADIENT_MODE = True
+                else:
+                    logger.warning(f"🔄 Diffusers import issue: {e}")
+    except Exception as e:
+        logger.warning(f"🔄 PyTorch/Diffusers compatibility check failed: {e}")
+        FORCE_GRADIENT_MODE = True
+else:
+    logger.info("ℹ️  Skipping diffusers check - PyTorch not available")
+
+# Override diffusers availability if we're forcing gradient mode
+if FORCE_GRADIENT_MODE:
     DIFFUSERS_AVAILABLE = False
 
 class AIBackgroundGenerator:
     """
-    AI Background Generator with graceful fallbacks
+    AI Background Generator with intelligent gradient fallbacks
     """
     
     def __init__(self):
         self.pipeline = None
         self.device = "cpu"
         
-        # Color themes for gradient fallbacks
+        # Comprehensive color themes for gradient fallbacks
         self.color_themes = {
+            # Blues
             'blue': [(64, 128, 255), (0, 64, 128)],
-            'ocean': [(0, 119, 190), (0, 64, 128)],
+            'ocean': [(0, 119, 190), (0, 64, 128)], 
             'sky': [(135, 206, 250), (25, 25, 112)],
+            'water': [(0, 191, 255), (0, 100, 200)],
+            'azure': [(240, 255, 255), (0, 127, 255)],
+            
+            # Greens
             'green': [(34, 139, 34), (0, 100, 0)],
             'nature': [(107, 142, 35), (34, 139, 34)],
             'forest': [(34, 139, 34), (0, 50, 0)],
+            'grass': [(124, 252, 0), (34, 139, 34)],
+            'mint': [(152, 251, 152), (0, 128, 0)],
+            
+            # Professional/Business
             'professional': [(105, 105, 105), (169, 169, 169)],
             'office': [(192, 192, 192), (105, 105, 105)],
             'corporate': [(70, 130, 180), (25, 25, 112)],
+            'business': [(47, 79, 79), (112, 128, 144)],
+            'modern': [(95, 158, 160), (47, 79, 79)],
+            
+            # Dark themes
             'dark': [(64, 64, 64), (0, 0, 0)],
             'night': [(25, 25, 112), (0, 0, 0)],
+            'black': [(105, 105, 105), (0, 0, 0)],
+            'shadow': [(85, 85, 85), (0, 0, 0)],
+            
+            # Warm colors
             'warm': [(255, 140, 0), (255, 69, 0)],
             'sunset': [(255, 94, 77), (255, 154, 0)],
+            'orange': [(255, 165, 0), (255, 69, 0)],
+            'fire': [(255, 69, 0), (139, 0, 0)],
+            
+            # Cool colors
             'purple': [(147, 112, 219), (75, 0, 130)],
+            'violet': [(138, 43, 226), (75, 0, 130)],
+            'lavender': [(230, 230, 250), (147, 112, 219)],
+            
+            # Others
             'red': [(220, 20, 60), (139, 0, 0)],
+            'pink': [(255, 182, 193), (255, 20, 147)],
+            'yellow': [(255, 255, 0), (255, 215, 0)],
+            'gold': [(255, 215, 0), (184, 134, 11)],
+            
+            # Technology/Digital
+            'tech': [(0, 255, 255), (0, 0, 139)],
+            'digital': [(138, 43, 226), (25, 25, 112)],
+            'cyber': [(0, 255, 127), (0, 100, 0)],
+            'neon': [(255, 20, 147), (138, 43, 226)],
+            
+            # Default
             'default': [(100, 149, 237), (65, 105, 225)]
         }
         
-        if DIFFUSERS_AVAILABLE:
+        # Only try to initialize diffusers if it's actually available
+        if DIFFUSERS_AVAILABLE and not FORCE_GRADIENT_MODE:
+            logger.info("🎨 Attempting to initialize AI pipeline...")
             self._init_diffusers()
+        else:
+            logger.info("🎨 Using gradient-only mode")
     
     def _init_diffusers(self):
-        """Initialize the Stable Diffusion pipeline"""
+        """Initialize the Stable Diffusion pipeline (only if safe)"""
+        # This will only be called if we've verified diffusers works
         try:
-            # Use a lightweight model for faster generation
+            from diffusers import StableDiffusionPipeline
+            
             model_id = "runwayml/stable-diffusion-v1-5"
             
             if torch.cuda.is_available():
                 self.device = "cuda"
-                logger.info("Using CUDA device")
+                logger.info("🚀 Using CUDA device")
             else:
                 self.device = "cpu"
-                logger.info("Using CPU device")
+                logger.info("🖥️  Using CPU device")
             
             self.pipeline = StableDiffusionPipeline.from_pretrained(
                 model_id,
@@ -95,10 +187,10 @@ def _init_diffusers(self):
                 self.pipeline.enable_memory_efficient_attention()
                 self.pipeline.enable_attention_slicing()
             
-            logger.info("Diffusers pipeline initialized successfully")
+            logger.info("✅ AI pipeline initialized successfully")
             
         except Exception as e:
-            logger.error(f"Failed to initialize diffusers pipeline: {e}")
+            logger.error(f"❌ Failed to initialize AI pipeline: {e}")
             self.pipeline = None
             global DIFFUSERS_AVAILABLE
             DIFFUSERS_AVAILABLE = False
@@ -107,28 +199,51 @@ def _analyze_prompt_theme(self, prompt: str) -> str:
         """Analyze prompt to determine appropriate color theme"""
         prompt_lower = prompt.lower()
         
-        # Check for specific themes
-        for theme, colors in self.color_themes.items():
+        # Direct theme matches first
+        for theme in self.color_themes:
             if theme in prompt_lower:
                 return theme
         
-        # Check for related keywords
-        if any(word in prompt_lower for word in ['water', 'sea', 'lake', 'river']):
-            return 'blue'
-        elif any(word in prompt_lower for word in ['grass', 'tree', 'plant', 'garden']):
-            return 'green'
-        elif any(word in prompt_lower for word in ['business', 'meeting', 'work']):
-            return 'professional'
-        elif any(word in prompt_lower for word in ['evening', 'midnight', 'shadow']):
-            return 'dark'
-        elif any(word in prompt_lower for word in ['fire', 'flame', 'autumn']):
-            return 'warm'
-        else:
-            return 'default'
+        # Keyword analysis
+        keyword_map = {
+            # Water/Ocean
+            ('water', 'sea', 'lake', 'river', 'stream'): 'ocean',
+            ('sky', 'cloud', 'air'): 'sky',
+            
+            # Nature
+            ('tree', 'plant', 'garden', 'leaf', 'flower'): 'nature',
+            ('grass', 'field', 'meadow'): 'grass',
+            
+            # Business
+            ('business', 'meeting', 'work', 'conference'): 'business',
+            ('company', 'enterprise', 'corporate'): 'corporate',
+            
+            # Time/Lighting
+            ('evening', 'midnight', 'shadow', 'darkness'): 'night',
+            ('fire', 'flame', 'autumn', 'fall'): 'fire',
+            ('morning', 'sunrise', 'dawn'): 'warm',
+            
+            # Technology
+            ('technology', 'computer', 'digital', 'software'): 'tech',
+            ('cyber', 'virtual', 'matrix'): 'cyber',
+            ('neon', 'electric', 'bright'): 'neon',
+            
+            # Emotions/Moods
+            ('calm', 'peaceful', 'serene'): 'azure',
+            ('energetic', 'vibrant', 'active'): 'orange',
+            ('elegant', 'sophisticated', 'luxury'): 'purple',
+            ('fresh', 'clean', 'pure'): 'mint',
+        }
+        
+        for keywords, theme in keyword_map.items():
+            if any(keyword in prompt_lower for keyword in keywords):
+                return theme
+        
+        return 'default'
     
     def _create_gradient_background(self, width: int = 1024, height: int = 768, 
                                   theme: str = 'default') -> Image.Image:
-        """Create a gradient background as fallback"""
+        """Create a sophisticated gradient background"""
         if not PIL_AVAILABLE:
             raise RuntimeError("PIL is not available for gradient generation")
         
@@ -136,14 +251,16 @@ def _create_gradient_background(self, width: int = 1024, height: int = 768,
         colors = self.color_themes.get(theme, self.color_themes['default'])
         color1, color2 = colors
         
-        # Create gradient
+        # Create base gradient
         image = Image.new('RGB', (width, height))
         draw = ImageDraw.Draw(image)
         
-        # Create vertical gradient
+        # Create smooth vertical gradient
         for y in range(height):
-            # Calculate interpolation factor
+            # Use smooth interpolation
             factor = y / height
+            # Apply easing function for smoother gradients
+            factor = factor * factor * (3.0 - 2.0 * factor)  # Smoothstep
             
             # Interpolate between colors
             r = int(color1[0] * (1 - factor) + color2[0] * factor)
@@ -153,37 +270,69 @@ def _create_gradient_background(self, width: int = 1024, height: int = 768,
             # Draw horizontal line
             draw.line([(0, y), (width, y)], fill=(r, g, b))
         
-        # Add subtle noise for texture
-        noise_overlay = Image.new('RGB', (width, height))
-        noise_draw = ImageDraw.Draw(noise_overlay)
+        # Add subtle texture and depth
+        self._add_texture(image, theme)
         
-        for _ in range(width * height // 100):  # Sparse noise
-            x = random.randint(0, width - 1)
-            y = random.randint(0, height - 1)
-            brightness = random.randint(-20, 20)
-            noise_draw.point((x, y), fill=(brightness, brightness, brightness))
+        return image
+    
+    def _add_texture(self, image: Image.Image, theme: str):
+        """Add subtle texture to make gradients more interesting"""
+        width, height = image.size
         
-        # Apply noise with low opacity
-        image = Image.blend(image, noise_overlay, 0.1)
+        # Create texture overlay
+        texture = Image.new('RGBA', (width, height), (0, 0, 0, 0))
+        texture_draw = ImageDraw.Draw(texture)
         
-        # Apply slight blur for smoothness
-        image = image.filter(ImageFilter.GaussianBlur(radius=1))
+        # Add different textures based on theme
+        if theme in ['tech', 'digital', 'cyber']:
+            # Add subtle grid pattern
+            grid_size = 50
+            for x in range(0, width, grid_size):
+                texture_draw.line([(x, 0), (x, height)], fill=(255, 255, 255, 5))
+            for y in range(0, height, grid_size):
+                texture_draw.line([(0, y), (width, y)], fill=(255, 255, 255, 5))
         
-        return image
+        elif theme in ['nature', 'forest', 'grass']:
+            # Add organic noise
+            for _ in range(width * height // 200):
+                x = random.randint(0, width - 1)
+                y = random.randint(0, height - 1)
+                size = random.randint(1, 3)
+                brightness = random.randint(10, 30)
+                texture_draw.ellipse([(x, y), (x+size, y+size)], 
+                                   fill=(brightness, brightness, brightness, 20))
+        
+        else:
+            # Add subtle noise for general texture
+            for _ in range(width * height // 300):
+                x = random.randint(0, width - 1)
+                y = random.randint(0, height - 1)
+                brightness = random.randint(-15, 15)
+                alpha = random.randint(5, 15)
+                texture_draw.point((x, y), fill=(brightness, brightness, brightness, alpha))
+        
+        # Apply texture
+        image.paste(texture, (0, 0), texture)
+        
+        # Final smooth blur for professional look
+        image = image.filter(ImageFilter.GaussianBlur(radius=0.8))
     
     def generate_background(self, prompt: str, width: int = 1024, height: int = 768,
                           guidance_scale: float = 7.5, num_inference_steps: int = 20) -> Optional[Image.Image]:
         """
         Generate a background image from a text prompt.
-        Falls back to gradient if AI generation fails.
+        Uses gradients with intelligent theming.
         """
         
-        # First try AI generation if available
-        if DIFFUSERS_AVAILABLE and self.pipeline is not None:
+        if not PIL_AVAILABLE:
+            logger.error("❌ Cannot generate backgrounds - PIL not available")
+            return None
+        
+        # Check if AI generation is possible and working
+        if DIFFUSERS_AVAILABLE and not FORCE_GRADIENT_MODE and self.pipeline is not None:
             try:
-                logger.info(f"Generating AI background for prompt: '{prompt}'")
+                logger.info(f"🎨 Attempting AI generation for: '{prompt}'")
                 
-                # Enhance prompt for better backgrounds
                 enhanced_prompt = f"{prompt}, high quality, detailed, professional background, 8k"
                 
                 with torch.no_grad():
@@ -197,36 +346,42 @@ def generate_background(self, prompt: str, width: int = 1024, height: int = 768,
                     )
                 
                 if result.images and len(result.images) > 0:
-                    logger.info("AI background generated successfully")
+                    logger.info("✅ AI background generated successfully")
                     return result.images[0]
                 
             except Exception as e:
-                logger.error(f"AI generation failed: {e}")
+                logger.warning(f"⚠️  AI generation failed, using gradient: {e}")
         
-        # Fallback to gradient generation
-        logger.info("Using gradient fallback for background generation")
+        # Use gradient generation (main path for compatibility)
+        logger.info(f"🎨 Creating gradient background for: '{prompt}'")
         theme = self._analyze_prompt_theme(prompt)
-        logger.info(f"Selected theme '{theme}' for prompt: '{prompt}'")
+        logger.info(f"🎯 Selected theme: '{theme}'")
         
         try:
-            return self._create_gradient_background(width, height, theme)
+            image = self._create_gradient_background(width, height, theme)
+            if image:
+                logger.info("✅ Gradient background generated successfully")
+                return image
+            else:
+                logger.error("❌ Gradient generation returned None")
+                return None
         except Exception as e:
-            logger.error(f"Gradient generation failed: {e}")
+            logger.error(f"❌ Gradient generation failed: {e}")
             return None
     
     def save_background(self, image: Image.Image, output_path: str) -> bool:
         """Save the generated background image"""
         try:
             # Ensure directory exists
-            os.makedirs(os.path.dirname(output_path), exist_ok=True)
+            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
             
-            # Save image
-            image.save(output_path, format='PNG', quality=95)
-            logger.info(f"Background saved to: {output_path}")
+            # Save image with high quality
+            image.save(output_path, format='PNG', quality=95, optimize=True)
+            logger.info(f"💾 Background saved to: {output_path}")
             return True
             
         except Exception as e:
-            logger.error(f"Failed to save background: {e}")
+            logger.error(f"❌ Failed to save background: {e}")
             return False
     
     def get_background_base64(self, image: Image.Image) -> str:
@@ -237,10 +392,10 @@ def get_background_base64(self, image: Image.Image) -> str:
             img_str = base64.b64encode(buffer.getvalue()).decode()
             return img_str
         except Exception as e:
-            logger.error(f"Failed to convert to base64: {e}")
+            logger.error(f"❌ Failed to convert to base64: {e}")
             return ""
 
-# Convenience function for quick background generation
+# Convenience functions
 def generate_ai_background(prompt: str, width: int = 1024, height: int = 768) -> Optional[Image.Image]:
     """
     Quick function to generate a background.
@@ -253,52 +408,90 @@ def generate_ai_background(prompt: str, width: int = 1024, height: int = 768) ->
     Returns:
         PIL Image object or None if generation fails
     """
-    generator = AIBackgroundGenerator()
-    return generator.generate_background(prompt, width, height)
+    try:
+        generator = AIBackgroundGenerator()
+        return generator.generate_background(prompt, width, height)
+    except Exception as e:
+        logger.error(f"❌ Background generation failed: {e}")
+        return None
 
-# Example usage and testing
-if __name__ == "__main__":
-    # Quick test first
-    if not test_background_generation():
-        print("Exiting due to test failure")
-        sys.exit(1)
+def test_background_generation():
+    """Test function to verify the background generator is working"""
+    print("\n" + "="*60)
+    print("🧪 AI BACKGROUND GENERATOR COMPATIBILITY TEST")
+    print("="*60)
+    
+    print(f"📦 PIL Available: {'✅' if PIL_AVAILABLE else '❌'}")
+    print(f"🔥 PyTorch Available: {'✅' if TORCH_AVAILABLE else '❌'}")
+    print(f"🎨 Diffusers Available: {'✅' if DIFFUSERS_AVAILABLE else '❌'}")
+    print(f"🔄 Force Gradient Mode: {'✅' if FORCE_GRADIENT_MODE else '❌'}")
     
-    # Full test suite
-    test_prompts = [
-        "professional office environment",
-        "peaceful ocean sunset", 
-        "dark forest at night",
-        "modern corporate meeting room",
-        "abstract blue technology background"
+    if not PIL_AVAILABLE:
+        print("\n❌ CRITICAL: Cannot generate backgrounds - PIL not available")
+        return False
+    
+    mode = "AI (Stable Diffusion)" if (DIFFUSERS_AVAILABLE and not FORCE_GRADIENT_MODE) else "Gradient Fallback"
+    print(f"\n🎯 Generation Mode: {mode}")
+    
+    # Test multiple themes
+    test_cases = [
+        ("professional blue office", "Should produce blue professional gradient"),
+        ("ocean sunset background", "Should produce ocean-themed gradient"),
+        ("dark tech cyber background", "Should produce dark tech gradient"),
+        ("green nature forest", "Should produce green nature gradient")
     ]
     
-    generator = AIBackgroundGenerator()
+    print(f"\n🔄 Testing {len(test_cases)} different prompts...")
     
-    for i, prompt in enumerate(test_prompts):
-        print(f"\nGenerating background {i+1}: '{prompt}'")
+    success_count = 0
+    for i, (prompt, expected) in enumerate(test_cases, 1):
+        print(f"\n📋 Test {i}: '{prompt}'")
+        print(f"    Expected: {expected}")
         
         try:
-            image = generator.generate_background(prompt, width=512, height=384)
+            image = generate_ai_background(prompt, 400, 300)
             
             if image:
-                output_path = f"test_background_{i+1}.png"
-                if generator.save_background(image, output_path):
-                    print(f"✓ Successfully saved: {output_path}")
-                else:
-                    print(f"✗ Failed to save: {output_path}")
+                print(f"    ✅ Generated: {image.size} {image.mode} image")
+                success_count += 1
+                
+                # Save test image
+                test_path = f"test_bg_{i}.png"
+                try:
+                    image.save(test_path)
+                    print(f"    💾 Saved: {test_path}")
+                except Exception as save_error:
+                    print(f"    ⚠️  Save failed: {save_error}")
             else:
-                print("✗ Failed to generate background")
+                print("    ❌ Generation failed - returned None")
                 
         except Exception as e:
-            print(f"✗ Error: {e}")
+            print(f"    ❌ Generation error: {e}")
+    
+    print(f"\n" + "="*60)
+    print(f"📊 RESULTS: {success_count}/{len(test_cases)} tests passed")
+    
+    if success_count == len(test_cases):
+        print("🎉 ALL TESTS PASSED! Background generator is working perfectly.")
+        return True
+    elif success_count > 0:
+        print("⚠️  PARTIAL SUCCESS: Some backgrounds generated successfully.")
+        return True
+    else:
+        print("❌ ALL TESTS FAILED: Background generator is not working.")
+        return False
+
+# Example usage and testing
+if __name__ == "__main__":
+    # Run compatibility test
+    success = test_background_generation()
     
-    print(f"\nFinal module status:")
-    print(f"- PIL Available: {PIL_AVAILABLE}")
-    print(f"- PyTorch Available: {TORCH_AVAILABLE}")
-    print(f"- Diffusers Available: {DIFFUSERS_AVAILABLE}")
-    if generator.pipeline:
-        print(f"- Device: {generator.device}")
-        print("- Mode: AI Generation (Stable Diffusion)")
+    if success:
+        print(f"\n🚀 Ready to generate backgrounds!")
+        print(f"💡 Usage example:")
+        print(f"   from ai_background import generate_ai_background")
+        print(f"   image = generate_ai_background('your prompt here')")
+        print(f"   image.save('background.png')")
     else:
-        print("- AI Pipeline: Not initialized")
-        print("- Mode: Gradient Fallback Only") 
+        print(f"\n⚠️  Please check the error messages above.")
+        print(f"💡 Make sure PIL (Pillow) is installed: pip install Pillow")