Update processing/ai_background.py

processing/ai_background.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """
-AI Background Generator Module
-Handles Stable Diffusion background generation with proper dependency management.
+AI Background Generator Module - Updated to handle dependency conflicts
+Handles Stable Diffusion background generation with graceful fallbacks.
 """
 
 import os
@@ -11,18 +11,21 @@
 import logging
 from pathlib import Path
 from typing import Optional
+import numpy as np
+import cv2
 
 logger = logging.getLogger(__name__)
 
 class AIBackgroundGenerator:
     """
-    Stable Diffusion background generator with dependency isolation.
+    AI background generator with dependency isolation and themed fallbacks.
     """
     
     def __init__(self, temp_dir: Optional[str] = None):
         self.temp_dir = temp_dir or tempfile.gettempdir()
         self.available = False
         self.error_message = None
+        self.use_fallback = False
         self._check_dependencies()
     
     def _check_dependencies(self):
@@ -30,36 +33,78 @@ def _check_dependencies(self):
         try:
             import torch
             if not torch.cuda.is_available():
-                self.error_message = "CUDA not available - AI background generation requires GPU"
+                self.error_message = "CUDA not available - using themed gradient fallback"
+                self.use_fallback = True
                 return
             
             # Test diffusers import
             from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
             
             self.available = True
-            logger.info("AI Background Generator: Dependencies OK")
+            logger.info("AI Background Generator: Full Stable Diffusion available")
             
         except ImportError as e:
             if "torch.library" in str(e) and "custom_op" in str(e):
-                self.error_message = (
-                    "PyTorch/Diffusers version mismatch. Please update:\n"
-                    "pip install --upgrade torch diffusers transformers accelerate"
-                )
+                self.error_message = "PyTorch/Diffusers version mismatch - using themed gradient fallback"
+                logger.warning(f"Diffusers unavailable due to dependency conflict: {e}")
             else:
-                self.error_message = f"Missing dependencies: {e}"
-            logger.warning(f"AI Background Generator unavailable: {self.error_message}")
+                self.error_message = f"Missing dependencies - using themed gradient fallback: {e}"
+                logger.warning(f"AI Background Generator using fallback: {self.error_message}")
+            self.use_fallback = True
         except Exception as e:
-            self.error_message = f"Unexpected error checking dependencies: {e}"
+            self.error_message = f"Unexpected error - using themed gradient fallback: {e}"
             logger.error(f"AI Background Generator error: {self.error_message}")
+            self.use_fallback = True
     
     def is_available(self) -> bool:
-        """Check if AI background generation is available."""
-        return self.available
+        """Check if AI background generation is available (includes fallback)."""
+        return self.available or self.use_fallback
     
     def get_error_message(self) -> Optional[str]:
         """Get error message if dependencies are not available."""
         return self.error_message
     
+    def _make_vertical_gradient(self, width: int, height: int, c1: tuple, c2: tuple) -> np.ndarray:
+        """Create a vertical gradient image."""
+        top = np.array(c1, dtype=np.float32)
+        bot = np.array(c2, dtype=np.float32)
+        rows = np.linspace(top, bot, num=max(1, height), dtype=np.float32)
+        grad = np.repeat(rows[:, None, :], repeats=max(1, width), axis=1)
+        return np.clip(grad, 0, 255).astype(np.uint8)
+    
+    def _generate_themed_gradient(self, width: int, height: int, prompt: str, seed: Optional[int] = None) -> str:
+        """Generate themed gradient based on prompt analysis."""
+        prompt_lower = prompt.lower()
+        
+        # Analyze prompt for appropriate colors
+        if any(word in prompt_lower for word in ["blue", "ocean", "sky", "water", "corporate"]):
+            colors = ((173, 216, 230), (0, 191, 255))  # Blue gradient
+        elif any(word in prompt_lower for word in ["green", "nature", "forest", "plant", "garden"]):
+            colors = ((34, 139, 34), (144, 238, 144))  # Green gradient
+        elif any(word in prompt_lower for word in ["warm", "office", "professional", "beige", "neutral"]):
+            colors = ((240, 240, 245), (200, 200, 210))  # Warm office
+        elif any(word in prompt_lower for word in ["dark", "night", "black", "studio", "dramatic"]):
+            colors = ((20, 20, 25), (40, 40, 45))  # Dark professional
+        elif any(word in prompt_lower for word in ["sunset", "orange", "pink", "warm"]):
+            colors = ((255, 182, 193), (255, 105, 180))  # Sunset pink
+        elif any(word in prompt_lower for word in ["purple", "violet", "magenta"]):
+            colors = ((147, 112, 219), (216, 191, 216))  # Purple gradient
+        elif any(word in prompt_lower for word in ["red", "burgundy", "wine"]):
+            colors = ((220, 20, 60), (255, 182, 193))  # Red gradient
+        else:
+            colors = ((40, 40, 48), (96, 96, 112))  # Default slate
+        
+        # Create gradient
+        grad = self._make_vertical_gradient(width, height, colors[0], colors[1])
+        
+        # Save result
+        seed_str = f"_{seed:08x}" if seed else ""
+        output_path = os.path.join(self.temp_dir, f"themed_bg_{int(time.time())}{seed_str}.jpg")
+        cv2.imwrite(output_path, grad)
+        
+        logger.info(f"Generated themed gradient for prompt: '{prompt}' -> {colors}")
+        return output_path
+    
     def generate_background(
         self,
         width: int,
@@ -73,22 +118,27 @@ def generate_background(
         seed: Optional[int] = None,
     ) -> str:
         """
-        Generate AI background image.
+        Generate AI background image with fallback to themed gradients.
         
         Returns:
             Path to generated background image
             
         Raises:
-            RuntimeError: If dependencies are not available or generation fails
+            RuntimeError: If both AI generation and fallback fail
         """
-        if not self.available:
-            raise RuntimeError(f"AI Background not available: {self.error_message}")
         
+        # Use fallback if AI generation unavailable
+        if self.use_fallback or not self.available:
+            logger.info(f"Using themed gradient fallback for: '{prompt}'")
+            return self._generate_themed_gradient(width, height, prompt, seed)
+        
+        # Try full AI generation
         try:
             # Import here to avoid issues if dependencies not available
             import torch
             from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
             from PIL import Image
+            import time
             
             device = "cuda" if torch.cuda.is_available() else "cpu"
             torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
@@ -99,7 +149,7 @@ def generate_background(
                 seed = random.randint(0, 2**31 - 1)
             generator.manual_seed(seed)
             
-            logger.info(f"Generating {width}x{height} background: '{prompt}' (seed: {seed})")
+            logger.info(f"Generating {width}x{height} AI background: '{prompt}' (seed: {seed})")
             
             # Choose pipeline based on whether we have an init image
             if init_image_path and os.path.exists(init_image_path):
@@ -165,7 +215,7 @@ def generate_background(
             # Save result
             output_path = os.path.join(
                 self.temp_dir, 
-                f"ai_bg_{int(os.times().elapsed)}_{seed:08x}.jpg"
+                f"ai_bg_{int(time.time())}_{seed:08x}.jpg"
             )
             result.save(output_path, quality=95, optimize=True)
             
@@ -176,12 +226,13 @@ def generate_background(
             except Exception:
                 pass
             
-            logger.info(f"AI background generated: {output_path}")
+            logger.info(f"AI background generated successfully: {output_path}")
             return output_path
             
         except Exception as e:
-            logger.error(f"AI background generation failed: {e}")
-            raise RuntimeError(f"Background generation failed: {e}")
+            logger.warning(f"AI generation failed: {e}, falling back to themed gradient")
+            # Fall back to themed gradient on any error
+            return self._generate_themed_gradient(width, height, prompt, seed)
 
 
 # Convenience function for easy import