Update generator.py

generator.py

@@ -1,6 +1,7 @@
 """
 Generation logic for Pixagram AI Pixel Art Generator
 """
+import gc
 import torch
 import numpy as np
 import cv2
@@ -196,55 +197,62 @@ class RetroArtConverter:
             print(f"[INFO] Verification skipped: {e}")
         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:
-                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)
-                    
-                    # 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))
-                    
-                    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}")
-                    
-                    # FIXED: Add torch.no_grad() wrapper
-                    with torch.no_grad():
-                        depth_image = self.depth_detector(image_for_depth)
-                    
-                    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")
-                    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-                    depth_colored = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
-                    return Image.fromarray(depth_colored)
-            else:
-                # No depth detector available, use grayscale fallback
-                print("[DEPTH] No depth detector available, using grayscale fallback")
+        """
+        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():
+                    depth_image = self.depth_detector(image_for_depth)
+                
+                # 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)    
+                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 add_trigger_word(self, prompt):