Create ai_background.py

processing/ai_background.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+"""
+AI Background Generator Module
+Handles Stable Diffusion background generation with proper dependency management.
+"""
+
+import os
+import sys
+import tempfile
+import random
+import logging
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+class AIBackgroundGenerator:
+    """
+    Stable Diffusion background generator with dependency isolation.
+    """
+    
+    def __init__(self, temp_dir: Optional[str] = None):
+        self.temp_dir = temp_dir or tempfile.gettempdir()
+        self.available = False
+        self.error_message = None
+        self._check_dependencies()
+    
+    def _check_dependencies(self):
+        """Check if required dependencies are available."""
+        try:
+            import torch
+            if not torch.cuda.is_available():
+                self.error_message = "CUDA not available - AI background generation requires GPU"
+                return
+            
+            # Test diffusers import
+            from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
+            
+            self.available = True
+            logger.info("AI Background Generator: Dependencies OK")
+            
+        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"
+                )
+            else:
+                self.error_message = f"Missing dependencies: {e}"
+            logger.warning(f"AI Background Generator unavailable: {self.error_message}")
+        except Exception as e:
+            self.error_message = f"Unexpected error checking dependencies: {e}"
+            logger.error(f"AI Background Generator error: {self.error_message}")
+    
+    def is_available(self) -> bool:
+        """Check if AI background generation is available."""
+        return self.available
+    
+    def get_error_message(self) -> Optional[str]:
+        """Get error message if dependencies are not available."""
+        return self.error_message
+    
+    def generate_background(
+        self,
+        width: int,
+        height: int,
+        prompt: str,
+        init_image_path: Optional[str] = None,
+        model_id: str = "runwayml/stable-diffusion-v1-5",
+        num_steps: int = 25,
+        guidance_scale: float = 7.5,
+        strength: float = 0.6,
+        seed: Optional[int] = None,
+    ) -> str:
+        """
+        Generate AI background image.
+        
+        Returns:
+            Path to generated background image
+            
+        Raises:
+            RuntimeError: If dependencies are not available or generation fails
+        """
+        if not self.available:
+            raise RuntimeError(f"AI Background not available: {self.error_message}")
+        
+        try:
+            # Import here to avoid issues if dependencies not available
+            import torch
+            from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
+            from PIL import Image
+            
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+            
+            # Setup generator
+            generator = torch.Generator(device=device)
+            if seed is None:
+                seed = random.randint(0, 2**31 - 1)
+            generator.manual_seed(seed)
+            
+            logger.info(f"Generating {width}x{height} 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):
+                # Image-to-image pipeline
+                logger.info("Using img2img pipeline")
+                pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=torch_dtype,
+                    safety_checker=None,
+                    requires_safety_checker=False
+                ).to(device)
+                
+                # Enable memory efficient attention if available
+                try:
+                    pipe.enable_attention_slicing()
+                    pipe.enable_model_cpu_offload()
+                except AttributeError:
+                    pass
+                
+                # Load and resize init image
+                init_image = Image.open(init_image_path).convert("RGB")
+                init_image = init_image.resize((width, height), Image.LANCZOS)
+                
+                # Generate
+                result = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    strength=strength,
+                    num_inference_steps=num_steps,
+                    guidance_scale=guidance_scale,
+                    generator=generator,
+                    height=height,
+                    width=width
+                ).images[0]
+                
+            else:
+                # Text-to-image pipeline
+                logger.info("Using txt2img pipeline")
+                pipe = StableDiffusionPipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=torch_dtype,
+                    safety_checker=None,
+                    requires_safety_checker=False
+                ).to(device)
+                
+                # Enable memory efficient attention if available
+                try:
+                    pipe.enable_attention_slicing()
+                    pipe.enable_model_cpu_offload()
+                except AttributeError:
+                    pass
+                
+                # Generate
+                result = pipe(
+                    prompt=prompt,
+                    height=height,
+                    width=width,
+                    num_inference_steps=num_steps,
+                    guidance_scale=guidance_scale,
+                    generator=generator
+                ).images[0]
+            
+            # Save result
+            output_path = os.path.join(
+                self.temp_dir, 
+                f"ai_bg_{int(os.times().elapsed)}_{seed:08x}.jpg"
+            )
+            result.save(output_path, quality=95, optimize=True)
+            
+            # Cleanup GPU memory
+            try:
+                del pipe
+                torch.cuda.empty_cache() if torch.cuda.is_available() else None
+            except Exception:
+                pass
+            
+            logger.info(f"AI background generated: {output_path}")
+            return output_path
+            
+        except Exception as e:
+            logger.error(f"AI background generation failed: {e}")
+            raise RuntimeError(f"Background generation failed: {e}")
+
+
+# Convenience function for easy import
+def create_ai_background_generator(temp_dir: Optional[str] = None) -> AIBackgroundGenerator:
+    """Factory function to create AI background generator."""
+    return AIBackgroundGenerator(temp_dir)
+
+
+# Test function
+def test_ai_background():
+    """Test AI background generation."""
+    generator = create_ai_background_generator()
+    
+    if not generator.is_available():
+        print(f"AI Background not available: {generator.get_error_message()}")
+        return False
+    
+    try:
+        # Test with simple prompt
+        bg_path = generator.generate_background(
+            width=512,
+            height=512,
+            prompt="professional office background with soft lighting",
+            num_steps=20
+        )
+        print(f"Test successful: {bg_path}")
+        return True
+    except Exception as e:
+        print(f"Test failed: {e}")
+        return False
+
+
+if __name__ == "__main__":
+    # Run test when executed directly
+    test_ai_background()