Update generator.py

generator.py

@@ -97,50 +97,65 @@ class RetroArtConverter:
         print("===================\n")
     
     def get_depth_map(self, image):
-        """Generate depth map using Zoe Depth"""
-        if self.zoe_depth is not None:
-            try:
-                if image.mode != 'RGB':
-                    image = image.convert('RGB')
-                
-                # Use safe helpers for type safety
-                orig_width, orig_height = safe_image_size(image)
-                
-                # FIXED: Use multiples of 64 (not 32)
-                target_width = ensure_int((orig_width // 64) * 64)
-                target_height = ensure_int((orig_height // 64) * 64)
-                
-                target_width = ensure_int(max(64, target_width))
-                target_height = ensure_int(max(64, target_height))
-                
-                # Create an explicit tuple of standard ints
-                size_for_depth = (target_width, target_height)
-                
-                # Always resize using the explicit int tuple
-                image_for_depth = image.resize(size_for_depth, Image.LANCZOS)
-                
-                # Generate depth map
-                depth_image = self.zoe_depth(image_for_depth, detect_resolution=512, image_resolution=1024)
-                
-                # Resize to match original if needed
-                if (depth_image.width, depth_image.height) != (orig_width, orig_height):
-                    depth_image = depth_image.resize((orig_width, orig_height), Image.LANCZOS)
-                
-                # Convert to RGB if needed
-                if depth_image.mode != 'RGB':
-                    depth_image = depth_image.convert('RGB')
-                
-                return depth_image, np.array(depth_image)
-                
-            except Exception as e:
-                print(f"Depth map generation failed: {e}")
-                import traceback
-                traceback.print_exc()
-                return None, None
-        else:
-            print("  Zoe Depth not available")
-            return None, None
-    
+            """
+            Generate depth map using available depth detector.
+            Supports: LeresDetector, ZoeDetector, or MidasDetector.
+            """
+            if self.depth_detector is not None:
+                try:
+                    if image.mode != 'RGB':
+                        image = image.convert('RGB')
+                    
+                    orig_width, orig_height = image.size
+                    orig_width = int(orig_width)
+                    orig_height = int(orig_height)
+                    
+                    target_width = int((orig_width // 64) * 64)
+                    target_height = int((orig_height // 64) * 64)
+                    
+                    target_width = int(max(64, target_width))
+                    target_height = int(max(64, target_height))
+                    
+                    size_for_depth = (int(target_width), int(target_height))
+                    image_for_depth = image.resize(size_for_depth, Image.LANCZOS)
+                    
+                    if target_width != orig_width or target_height != orig_height:
+                        print(f"[DEPTH] Resized for {self.depth_type.upper()}Detector: {orig_width}x{orig_height} -> {target_width}x{target_height}")
+                    
+                    # Use torch.no_grad() and clear cache
+                    with torch.no_grad():
+                        # --- FIX: Move model to GPU for inference and back to CPU ---
+                        self.depth_detector.to(self.device)
+                        depth_image = self.depth_detector(image_for_depth)
+                        self.depth_detector.to("cpu")
+                    
+                    # ADDED: Clear GPU cache after depth detection
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                    
+                    depth_width, depth_height = depth_image.size
+                    if depth_width != orig_width or depth_height != orig_height:
+                        depth_image = depth_image.resize((int(orig_width), int(orig_height)), Image.LANCZOS)
+                    
+                    print(f"[DEPTH] {self.depth_type.upper()} depth map generated: {orig_width}x{orig_height}")
+                    return depth_image
+                    
+                except Exception as e:
+                    print(f"[DEPTH] {self.depth_type.upper()}Detector failed ({e}), falling back to grayscale depth")
+                    # ADDED: Clear cache on error
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                    
+                    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
+                    depth_colored = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
+                    return Image.fromarray(depth_colored)
+            else:
+                print("[DEPTH] No depth detector available, using grayscale fallback")
+                gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
+                depth_colored = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
+                return Image.fromarray(depth_colored)
+
+
     def generate_caption(self, image):
         """Generate caption for image using loaded caption model"""
         if not self.caption_enabled or self.caption_model is None: