import torch import numpy as np import cv2 from PIL import Image import logging import gc import time import os from typing import Optional, Dict, Any, Callable import warnings warnings.filterwarnings("ignore") from diffusers import StableDiffusionXLPipeline, StableDiffusionXLInpaintPipeline, DPMSolverMultistepScheduler import open_clip from mask_generator import MaskGenerator from image_blender import ImageBlender try: import spaces SPACES_AVAILABLE = True except ImportError: SPACES_AVAILABLE = False logger = logging.getLogger(__name__) class BackgroundEngine: """ Background generation engine for VividFlow. Integrates SDXL pipeline, OpenCLIP analysis, mask generation, and advanced image blending. """ def __init__(self, device: str = "auto"): self.device = self._setup_device(device) self.base_model_id = "stabilityai/stable-diffusion-xl-base-1.0" self.clip_model_name = "ViT-B-32" self.clip_pretrained = "openai" self.pipeline = None self.inpaint_pipeline = None self.clip_model = None self.clip_preprocess = None self.clip_tokenizer = None self.is_initialized = False self.inpaint_initialized = False self.max_image_size = 1024 self.default_steps = 25 self.use_fp16 = True self.mask_generator = MaskGenerator(self.max_image_size) self.image_blender = ImageBlender() logger.info(f"BackgroundEngine initialized on {self.device}") def _setup_device(self, device: str) -> str: """Setup computation device (ZeroGPU compatible)""" if os.getenv('SPACE_ID') is not None: return "cpu" if device == "auto": if torch.cuda.is_available(): return "cuda" elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): return "mps" return "cpu" return device def _memory_cleanup(self): """Memory cleanup""" for _ in range(3): gc.collect() is_spaces = os.getenv('SPACE_ID') is not None if not is_spaces and torch.cuda.is_available(): torch.cuda.empty_cache() def load_models(self, progress_callback: Optional[Callable] = None): """Load SDXL and OpenCLIP models""" if self.is_initialized: logger.info("Models already loaded") return logger.info("Loading background generation models...") try: self._memory_cleanup() # Detect actual device (in ZeroGPU, CUDA becomes available after @spaces.GPU allocation) actual_device = "cuda" if torch.cuda.is_available() else self.device logger.info(f"Loading models to device: {actual_device}") if progress_callback: progress_callback("Loading OpenCLIP...", 20) # Load OpenCLIP self.clip_model, _, self.clip_preprocess = open_clip.create_model_and_transforms( self.clip_model_name, pretrained=self.clip_pretrained, device=actual_device ) self.clip_tokenizer = open_clip.get_tokenizer(self.clip_model_name) self.clip_model.eval() logger.info("OpenCLIP loaded") if progress_callback: progress_callback("Loading SDXL pipeline...", 60) # Load SDXL self.pipeline = StableDiffusionXLPipeline.from_pretrained( self.base_model_id, torch_dtype=torch.float16 if self.use_fp16 else torch.float32, use_safetensors=True, variant="fp16" if self.use_fp16 else None ) # DPM solver for faster generation self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config( self.pipeline.scheduler.config ) self.pipeline = self.pipeline.to(actual_device) if progress_callback: progress_callback("Applying optimizations...", 90) # Memory optimizations try: self.pipeline.enable_xformers_memory_efficient_attention() logger.info("xformers enabled") except Exception: try: self.pipeline.enable_attention_slicing() logger.info("Attention slicing enabled") except Exception: pass if hasattr(self.pipeline, 'enable_vae_tiling'): self.pipeline.enable_vae_tiling() if hasattr(self.pipeline, 'enable_vae_slicing'): self.pipeline.enable_vae_slicing() self.pipeline.unet.eval() if hasattr(self.pipeline, 'vae'): self.pipeline.vae.eval() self.is_initialized = True if progress_callback: progress_callback("Models loaded!", 100) logger.info("Background models loaded successfully") except Exception as e: logger.error(f"Model loading failed: {e}") raise RuntimeError(f"Failed to load models: {str(e)}") def analyze_image_with_clip(self, image: Image.Image) -> str: """Analyze image using OpenCLIP""" if not self.clip_model: return "Unknown" try: # Use actual device actual_device = "cuda" if torch.cuda.is_available() else self.device image_input = self.clip_preprocess(image).unsqueeze(0).to(actual_device) categories = [ "a photo of a person", "a photo of an animal", "a photo of an object", "a photo of nature", "a photo of a building" ] text_inputs = self.clip_tokenizer(categories).to(actual_device) with torch.no_grad(): image_features = self.clip_model.encode_image(image_input) text_features = self.clip_model.encode_text(text_inputs) image_features /= image_features.norm(dim=-1, keepdim=True) text_features /= text_features.norm(dim=-1, keepdim=True) similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1) best_match_idx = similarity.argmax().item() category = categories[best_match_idx].replace("a photo of ", "") return category except Exception as e: logger.error(f"CLIP analysis failed: {e}") return "unknown" def enhance_prompt(self, user_prompt: str, foreground_image: Image.Image) -> str: """Smart prompt enhancement based on image analysis""" try: img_array = np.array(foreground_image.convert('RGB')) # Analyze color temperature lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB) avg_b = np.mean(lab[:, :, 2]) is_warm = avg_b > 128 # Analyze brightness gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY) avg_brightness = np.mean(gray) is_bright = avg_brightness > 127 # Get subject type clip_analysis = self.analyze_image_with_clip(foreground_image) subject_type = clip_analysis # Build lighting descriptors if is_warm and is_bright: lighting = "warm golden hour lighting, soft natural light" elif is_warm and not is_bright: lighting = "warm ambient lighting, cozy atmosphere" elif not is_warm and is_bright: lighting = "bright daylight, clear sky lighting" else: lighting = "soft diffused light, gentle shadows" # Build atmosphere based on subject atmosphere_map = { "person": "professional, elegant composition", "animal": "natural, harmonious setting", "object": "clean product photography style", "nature": "scenic, peaceful atmosphere", "building": "architectural, balanced composition" } atmosphere = atmosphere_map.get(subject_type, "balanced composition") quality_modifiers = "high quality, detailed, sharp focus, photorealistic" # Avoid conflicts user_prompt_lower = user_prompt.lower() if "sunset" in user_prompt_lower or "golden" in user_prompt_lower: lighting = "" if "dark" in user_prompt_lower or "night" in user_prompt_lower: lighting = lighting.replace("bright", "").replace("daylight", "") # Combine fragments = [user_prompt] if lighting: fragments.append(lighting) fragments.append(atmosphere) fragments.append(quality_modifiers) enhanced_prompt = ", ".join(filter(None, fragments)) logger.debug(f"Enhanced: {enhanced_prompt[:80]}...") return enhanced_prompt except Exception as e: logger.warning(f"Prompt enhancement failed: {e}") return f"{user_prompt}, high quality, detailed, photorealistic" def _prepare_image(self, image: Image.Image) -> Image.Image: """Prepare image for processing""" if image.mode != 'RGB': image = image.convert('RGB') width, height = image.size max_size = self.max_image_size if width > max_size or height > max_size: ratio = min(max_size/width, max_size/height) new_width = int(width * ratio) new_height = int(height * ratio) image = image.resize((new_width, new_height), Image.LANCZOS) width, height = image.size new_width = (width // 8) * 8 new_height = (height // 8) * 8 if new_width != width or new_height != height: image = image.resize((new_width, new_height), Image.LANCZOS) return image def generate_background( self, prompt: str, width: int, height: int, negative_prompt: str = "blurry, low quality, distorted", num_inference_steps: int = 25, guidance_scale: float = 7.5 ) -> Image.Image: """Generate background using SDXL""" if not self.is_initialized: raise RuntimeError("Models not loaded") logger.info(f"Generating background: {prompt[:50]}...") try: # Use actual device actual_device = "cuda" if torch.cuda.is_available() else self.device with torch.inference_mode(): result = self.pipeline( prompt=prompt, negative_prompt=negative_prompt, width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=torch.Generator(device=actual_device).manual_seed(42) ) generated_image = result.images[0] logger.info("Background generation completed") return generated_image except torch.cuda.OutOfMemoryError: logger.error("GPU memory exhausted") self._memory_cleanup() raise RuntimeError("GPU memory insufficient") except Exception as e: logger.error(f"Generation failed: {e}") raise RuntimeError(f"Generation failed: {str(e)}") def generate_and_combine( self, original_image: Image.Image, prompt: str, combination_mode: str = "center", focus_mode: str = "person", negative_prompt: str = "blurry, low quality, distorted", num_inference_steps: int = 25, guidance_scale: float = 7.5, progress_callback: Optional[Callable] = None, enable_prompt_enhancement: bool = True, feather_radius: int = 0, enhance_dark_edges: bool = False ) -> Dict[str, Any]: """ Generate background and combine with foreground. Args: feather_radius: Gaussian blur radius for mask edge softening (0-20, default 0) enhance_dark_edges: Enhance mask edges for dark background images (default False) Returns dict with: combined_image, generated_scene, original_image, mask, success """ if not self.is_initialized: raise RuntimeError("Models not loaded") logger.info("Starting background generation and combination...") try: if progress_callback: progress_callback("Analyzing image...", 5) # Prepare image processed_original = self._prepare_image(original_image) target_width, target_height = processed_original.size if progress_callback: progress_callback("Enhancing prompt...", 15) # Enhance prompt if enable_prompt_enhancement: enhanced_prompt = self.enhance_prompt(prompt, processed_original) else: enhanced_prompt = f"{prompt}, high quality, detailed, photorealistic" enhanced_negative = f"{negative_prompt}, people, characters, cartoons, logos" if progress_callback: progress_callback("Generating background...", 30) # Generate background generated_background = self.generate_background( prompt=enhanced_prompt, width=target_width, height=target_height, negative_prompt=enhanced_negative, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale ) if progress_callback: progress_callback("Creating mask...", 80) # Generate mask logger.info("Generating mask...") combination_mask = self.mask_generator.create_gradient_based_mask( processed_original, combination_mode, focus_mode, enhance_dark_edges=enhance_dark_edges ) if progress_callback: progress_callback("Blending images...", 90) # Blend images with feather_radius logger.info("Blending images...") combined_image = self.image_blender.simple_blend_images( processed_original, generated_background, combination_mask, feather_radius=feather_radius ) # Cleanup self._memory_cleanup() if progress_callback: progress_callback("Complete!", 100) logger.info("Background generation completed successfully") # Build result dict (always include mask for diagnostics) return { "combined_image": combined_image, "generated_scene": generated_background, "original_image": processed_original, "mask": combination_mask, "success": True } except Exception as e: logger.error(f"Generation failed: {e}") self._memory_cleanup() return { "success": False, "error": str(e) } def _load_inpaint_pipeline(self) -> bool: """Lazy load SDXL inpainting pipeline""" if self.inpaint_initialized: return True try: logger.info("Loading SDXL inpainting pipeline...") actual_device = "cuda" if torch.cuda.is_available() else self.device self.inpaint_pipeline = StableDiffusionXLInpaintPipeline.from_pretrained( "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16 if actual_device == "cuda" else torch.float32, variant="fp16" if actual_device == "cuda" else None, use_safetensors=True ) self.inpaint_pipeline.to(actual_device) # Use fast scheduler self.inpaint_pipeline.scheduler = DPMSolverMultistepScheduler.from_config( self.inpaint_pipeline.scheduler.config ) # Memory optimization if actual_device == "cuda": try: self.inpaint_pipeline.enable_xformers_memory_efficient_attention() except Exception: pass self.inpaint_initialized = True logger.info("✓ SDXL inpainting pipeline loaded") return True except Exception as e: logger.error(f"Failed to load inpainting pipeline: {e}") self.inpaint_initialized = False return False def inpaint_region( self, image: Image.Image, mask: Image.Image, prompt: str, negative_prompt: str = "blurry, low quality, artifacts, seams", num_inference_steps: int = 20, guidance_scale: float = 7.5, strength: float = 0.99 ) -> Dict[str, Any]: """ Inpaint marked regions with background content. Args: image: The combined image with artifacts to fix mask: Binary mask where white = areas to inpaint prompt: Background description for inpainting negative_prompt: What to avoid num_inference_steps: Denoising steps (20 is usually enough) guidance_scale: How closely to follow prompt strength: How much to change masked area (0.99 = almost complete replacement) Returns: Dict with inpainted_image, success, error """ try: # Load inpainting pipeline if not already loaded if not self._load_inpaint_pipeline(): # Fallback to OpenCV inpainting return self._opencv_inpaint_fallback(image, mask) logger.info("Starting region inpainting...") # Prepare images image = self._prepare_image(image) mask = mask.resize(image.size, Image.LANCZOS).convert('L') # Ensure mask is properly binarized mask_array = np.array(mask) mask_array = (mask_array > 127).astype(np.uint8) * 255 mask = Image.fromarray(mask_array, mode='L') # Dilate mask slightly for better blending kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) mask_dilated = cv2.dilate(mask_array, kernel, iterations=1) mask = Image.fromarray(mask_dilated, mode='L') actual_device = "cuda" if torch.cuda.is_available() else self.device with torch.inference_mode(): result = self.inpaint_pipeline( prompt=prompt, negative_prompt=negative_prompt, image=image, mask_image=mask, width=image.size[0], height=image.size[1], num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, strength=strength, generator=torch.Generator(device=actual_device).manual_seed(42) ) inpainted = result.images[0] # Blend edges for smoother transition inpainted = self._blend_inpaint_edges(image, inpainted, mask) self._memory_cleanup() logger.info("✓ Region inpainting completed") return { "inpainted_image": inpainted, "success": True } except Exception as e: logger.error(f"Inpainting failed: {e}") self._memory_cleanup() return { "success": False, "error": str(e) } def _opencv_inpaint_fallback( self, image: Image.Image, mask: Image.Image ) -> Dict[str, Any]: """Fallback to OpenCV inpainting for small areas or when SDXL unavailable""" try: logger.info("Using OpenCV inpainting fallback...") img_array = np.array(image.convert('RGB')) mask_array = np.array(mask.convert('L')) # Binarize mask mask_binary = (mask_array > 127).astype(np.uint8) * 255 # Use Telea algorithm for natural results inpainted = cv2.inpaint( img_array, mask_binary, inpaintRadius=5, flags=cv2.INPAINT_TELEA ) result = Image.fromarray(inpainted) logger.info("✓ OpenCV inpainting completed") return { "inpainted_image": result, "success": True } except Exception as e: logger.error(f"OpenCV inpainting failed: {e}") return { "success": False, "error": str(e) } def _blend_inpaint_edges( self, original: Image.Image, inpainted: Image.Image, mask: Image.Image, feather_pixels: int = 8 ) -> Image.Image: """Blend inpainted region edges for seamless transition""" try: orig_array = np.array(original).astype(np.float32) inpaint_array = np.array(inpainted).astype(np.float32) mask_array = np.array(mask.convert('L')).astype(np.float32) / 255.0 # Create feathered mask for smooth blending if feather_pixels > 0: kernel_size = feather_pixels * 2 + 1 mask_feathered = cv2.GaussianBlur( mask_array, (kernel_size, kernel_size), feather_pixels / 2 ) else: mask_feathered = mask_array # Expand mask to 3 channels mask_3d = mask_feathered[:, :, np.newaxis] # Blend: inpainted in masked area, original elsewhere blended = inpaint_array * mask_3d + orig_array * (1 - mask_3d) blended = np.clip(blended, 0, 255).astype(np.uint8) return Image.fromarray(blended) except Exception as e: logger.warning(f"Edge blending failed: {e}, returning inpainted directly") return inpainted