Update generator.py

generator.py

@@ -21,6 +21,7 @@ from models import (
     load_sdxl_pipeline, load_lora, setup_compel,
     setup_scheduler, optimize_pipeline, load_caption_model, set_clip_skip
 )
+from memory_utils import MemoryManager, ModelOffloader
 
 
 class RetroArtConverter:
@@ -36,10 +37,13 @@ class RetroArtConverter:
             'zoe_depth': False
         }
         
-        # Load face analysis
+        # Initialize memory manager
+        self.memory_manager = MemoryManager(device=device, dtype=dtype, verbose=True)
+        
+        # Load face analysis (stays on CPU)
         self.face_app, self.face_detection_enabled = load_face_analysis()
         
-        # Load depth detector
+        # Load depth detector (starts on CPU)
         self.zoe_depth, zoe_success = load_depth_detector()
         self.models_loaded['zoe_depth'] = zoe_success
         
@@ -67,7 +71,7 @@ class RetroArtConverter:
         # Optimize
         optimize_pipeline(self.pipe)
         
-        # Load caption model
+        # Load caption model (starts on CPU)
         self.caption_processor, self.caption_model, self.caption_enabled, self.caption_model_type = load_caption_model()
         
         # Set CLIP skip
@@ -76,7 +80,10 @@ class RetroArtConverter:
         # Print status
         self._print_status()
         
-        print("  [OK] RetroArtConverter initialized with InstantID!")
+        # Initial memory cleanup
+        self.memory_manager.cleanup_memory(aggressive=True)
+        
+        print("  [OK] RetroArtConverter initialized with optimized memory management!")
     
     def _print_status(self):
         """Print model loading status"""
@@ -89,7 +96,7 @@ class RetroArtConverter:
         print("===================\n")
     
     def get_depth_map(self, image):
-        """Generate depth map using Zoe Depth"""
+        """Generate depth map using Zoe Depth with optimized GPU usage"""
         if self.zoe_depth is not None:
             try:
                 if image.mode != 'RGB':
@@ -108,16 +115,27 @@ class RetroArtConverter:
                 size_for_depth = (int(target_width), int(target_height))
                 image_for_depth = image.resize(size_for_depth, Image.LANCZOS)
                 
+                # Move depth model to GPU temporarily
+                self.zoe_depth = self.zoe_depth.to(self.device)
+                
+                # Generate depth map
                 depth_array = self.zoe_depth(image_for_depth, detect_resolution=512, image_resolution=1024)
                 depth_image = Image.fromarray(depth_array)
                 
+                # Move depth model back to CPU to free GPU memory
+                self.zoe_depth = self.zoe_depth.to("cpu")
+                torch.cuda.empty_cache() if torch.cuda.is_available() else None
+                
                 if depth_image.size != image.size:
                     depth_image = depth_image.resize(image.size, Image.LANCZOS)
                 
-                print(f"[DEPTH] Generated depth map: {depth_image.size}")
+                print(f"[DEPTH] Generated depth map: {depth_image.size} (model offloaded to CPU)")
                 return depth_image, depth_array
             except Exception as e:
                 print(f"[DEPTH] Generation failed: {e}, using grayscale")
+                # Ensure model is back on CPU even if error
+                if hasattr(self, 'zoe_depth') and self.zoe_depth is not None:
+                    self.zoe_depth = self.zoe_depth.to("cpu")
                 return image.convert('L').convert('RGB'), None
         else:
             print("[DEPTH] Detector not available, using grayscale")
@@ -162,11 +180,14 @@ class RetroArtConverter:
             return None
     
     def generate_caption(self, image):
-        """Generate caption for image"""
+        """Generate caption for image with optimized GPU usage"""
         if not self.caption_enabled or self.caption_model is None:
             return None
         
         try:
+            # Move caption model to GPU temporarily
+            self.caption_model = self.caption_model.to(self.device)
+            
             if self.caption_model_type == 'git':
                 inputs = self.caption_processor(images=image, return_tensors="pt").to(self.device)
                 generated_ids = self.caption_model.generate(**inputs, max_length=CAPTION_CONFIG['max_length'])
@@ -176,11 +197,19 @@ class RetroArtConverter:
                 generated_ids = self.caption_model.generate(**inputs, max_length=CAPTION_CONFIG['max_length'])
                 caption = self.caption_processor.decode(generated_ids[0], skip_special_tokens=True)
             else:
+                self.caption_model = self.caption_model.to("cpu")  # Move back to CPU
                 return None
             
+            # Move caption model back to CPU to free GPU memory
+            self.caption_model = self.caption_model.to("cpu")
+            torch.cuda.empty_cache() if torch.cuda.is_available() else None
+            
             return sanitize_text(caption)
         except Exception as e:
             print(f"[CAPTION] Generation failed: {e}")
+            # Ensure model is back on CPU even if error
+            if hasattr(self, 'caption_model') and self.caption_model is not None:
+                self.caption_model = self.caption_model.to("cpu")
             return None
     
     def generate_retro_art(
@@ -389,10 +418,26 @@ class RetroArtConverter:
             return generated_image
         
         finally:
-            # Memory cleanup
+            # Aggressive memory cleanup
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
-            gc.collect()
+                torch.cuda.synchronize()  # Ensure all GPU operations complete
+            
+            # Force garbage collection multiple times for thorough cleanup
+            for _ in range(3):
+                gc.collect()
+            
+            # Additional cleanup for large tensors
+            if 'pipe_kwargs' in locals():
+                for key in list(pipe_kwargs.keys()):
+                    if isinstance(pipe_kwargs.get(key), torch.Tensor):
+                        del pipe_kwargs[key]
+            
+            # Log memory status if in debug mode
+            if torch.cuda.is_available():
+                allocated = torch.cuda.memory_allocated() / 1024**3
+                reserved = torch.cuda.memory_reserved() / 1024**3
+                print(f"[MEMORY] GPU: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved")
 
 
 print("[OK] Generator class ready with InstantID support")