Update generator.py

generator.py

@@ -146,64 +146,56 @@ class RetroArtConverter:
         print("============================\n")
     
     def get_depth_map(self, image):
-            """
-            Generate depth map using available depth detector.
-            Supports: LeresDetector, ZoeDetector, or MidasDetector.
-            """
-            if self.depth_detector is not None:
+            """Generate depth map using Zoe Depth"""
+            if self.zoe_depth is not None:
                 try:
                     if image.mode != 'RGB':
                         image = image.convert('RGB')
                     
                     orig_width, orig_height = image.size
+                    # **FIX 1 START: Ensure all size variables are standard Python int**
                     orig_width = int(orig_width)
                     orig_height = int(orig_height)
                     
+                    # FIXED: Use multiples of 64 (not 32)
                     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))
                     
+                    # Create an explicit tuple of standard ints
                     size_for_depth = (int(target_width), int(target_height))
+                    
+                    # Always resize using the explicit int tuple to avoid numpy.int64 issues
+                    # This replaces the conditional resize
                     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}")
+                        print(f"[DEPTH] Resized for ZoeDetector: {orig_width}x{orig_height} -> {target_width}x{target_height}")
                     
-                    # Use torch.no_grad() and clear cache
+                    # FIXED: Add torch.no_grad() wrapper
                     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_image = self.zoe_depth(image_for_depth) # Use the correctly-typed resized image
                     
                     depth_width, depth_height = depth_image.size
                     if depth_width != orig_width or depth_height != orig_height:
+                        # Resize back to the original size that get_depth_map received
                         depth_image = depth_image.resize((int(orig_width), int(orig_height)), Image.LANCZOS)
+                    # **FIX 1 END**
                     
-                    print(f"[DEPTH] {self.depth_type.upper()} depth map generated: {orig_width}x{orig_height}")
+                    print(f"[DEPTH] Zoe 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()
-                    
+                    print(f"[DEPTH] ZoeDetector failed ({e}), falling back to grayscale depth")
                     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)
-
+                return Image.fromarray(depth_colored)    
 
     
     def add_trigger_word(self, prompt):